layer_to_df.py

# -*- coding: utf-8 -*-
"""
/***************************************************************************
 DataPlotlyDialog
                                 A QGIS plugin
 D3 Plots for QGIS
                             -------------------
        begin                : 2020-08-17
        git sha              : $Format:%H$
        copyright            : (C) 2017 by matteo ghetta
        email                : matteo.ghetta@gmail.com
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""

import pandas as pd
import numpy as np
from qgis.PyQt.QtCore import QVariant


def layer_to_dataframe(layer):
    """Converts a QgsVectorLayer in a pandas dataframe

    Args:
        layer (QgsVectorLayer): any type of QgsVectorLayer

    Returns:
        df [pandas dataframe]: the final pandas dataframe
    """

    # from QVariant type to numpy types
    qvariant_type_map = {
        1: np.bool, # boolean
        2: np.int, # int
        3: np.int, # unsigned int
        4: np.longlong, # longlong
        5: np.ulonglong, # ulonglong
        6: np.double, # double
        7: np.str, # character
        10: np.str, # string
        14: np.dtype('datetime64[D]'), # date
        15: np.dtype('datetime64[s]'), # time
        16: np.dtype('datetime64[s]'), # datetime
    }

    # create a dictionary with key from name and values as pd.Series and data type
    # e.g. d = {'so4': pd.Series([], dtype=np.int), 'admin': pd.Series([], dtype=np.str}
    d = {k.name(): pd.Series([], dtype=qvariant_type_map[k.type()]) for k in layer.fields()}

    # loop into the dictionary items and overwrite the dictionary with the 
    # pd.Series filled with attribute table values casted to the correct data type
    for k, v in d.items():
        for i in layer.getFeatures():
            d[k] = d[k].append(pd.Series(i[k]))

    # create the DataFrame
    df = pd.DataFrame(d)
    # set the DataFrame index from feature ids
    df.set_index([[i.id() for i in layer.getFeatures()]], inplace=True)

    return df