QGIS/python/plugins/processing/tools/_vector.py

import os
from PyQt4.QtCore import *
from qgis.core import *
from collections import Iterator

TYPE_MAP = {
    str: QVariant.String,
    float: QVariant.Double,
    int: QVariant.Int,
    bool: QVariant.Bool
}

GEOM_TYPE_MAP = {
    QGis.WKBPoint: 'Point',
    QGis.WKBLineString: 'LineString',
    QGis.WKBPolygon: 'Polygon',
    QGis.WKBMultiPoint: 'MultiPoint',
    QGis.WKBMultiLineString: 'MultiLineString',
    QGis.WKBMultiPolygon: 'MultiPolygon',
}

QGSGEOM = "qgs_geom"


def _toQgsGeometry(geom, geomtype):
    def toQgsPoint(p):
        if isinstance(p, tuple):
            return QgsPoint(p[0], p[1])
        elif isinstance(p, QgsPoint):
            return p
        else:
            return [toQgsPoint(item) for item in p]
    if isinstance(geom, QgsGeometry):
        return geom
    if geomtype == QGis.WKBPoint:
        return QgsGeometry.fromPoint(toQgsPoint(geom))
    elif geomtype == QGis.WKBLineString:
        return QgsGeometry.fromPolyline(toQgsPoint(geom))
    elif geomtype == QGis.WKBPolygon:
        return QgsGeometry.fromPolygon(toQgsPoint(geom))
    else:
        pass
        #TODO


class VectorLayer(object):

    def __init__(self, layer):
        self.layer = layer

    def __getattr__(self, name):
        return getattr(self.__dict__['layer'], name)

    def addFeature(self, *args):
        if len(args) == 2:
            feature = QgsFeature()
            feature.setGeometry(_toQgsGeometry(args[0], self.wkbType()))
            feature.setAttributes(args[1])
        if len(args) == 1:
            if isinstance(args[0], dict):
                feature = QgsFeature()
                attrs = [args[0].get(f.name(), None) for f in self.layer.dataProvider().fields()]
                feature.setAttributes(attrs)
                if QGSGEOM in args[0]:
                    feature.setGeometry[args[0][QGSGEOM]]
            elif isinstance(args[0], QgsFeature):
                feature = args[0]
        self.dataProvider().addFeatures([feature])

    def addField(self, name, fieldtype, defaultValue=None):
        fields = [f.name() for f in self.dataProvider().fields()]
        self.dataProvider().addAttributes([_toQgsField((name, fieldtype))])
        self.updateFields()
        if defaultValue is not None:
            idx = self.dataProvider().fieldNameIndex(name)
            for featureIdx, f in enumerate(self.getFeatures()):
                if callable(defaultValue):
                    attrs = {QGSGEOM: Geometry(f.geometry())}
                    for i, field in enumerate(fields):
                        attrs[field] = f.attributes()[i]
                        v = defaultValue(featureIdx, attrs)
                else:
                    v = defaultValue
                self.dataProvider().changeAttributeValues({f.id(): {idx: v}})

    def features(self):
        class todict(Iterator):

            def __init__(self, featureiter, fields):
                self.featureiter = featureiter
                self.fields = fields

            def next(self):
                feature = self.featureiter.next()

                class geomdict(dict):

                    def geom(self):
                        return self[QGSGEOM]
                d = geomdict()
                d[QGSGEOM] = Geometry(feature.geometry())
                for i, field in enumerate(self.fields):
                    d[field] = feature.attributes()[i]
                return d
        fields = [f.name() for f in self.dataProvider().fields()]
        return todict(self.getFeatures(), fields)

TYPE_MAP = {
    str: QVariant.String,
    float: QVariant.Double,
    int: QVariant.Int,
    bool: QVariant.Bool
}

GEOM_TYPE_MAP = {
    QGis.WKBPoint: 'Point',
    QGis.WKBLineString: 'LineString',
    QGis.WKBPolygon: 'Polygon',
    QGis.WKBMultiPoint: 'MultiPoint',
    QGis.WKBMultiLineString: 'MultiLineString',
    QGis.WKBMultiPolygon: 'MultiPolygon',
}


def _toQgsField(f):
    if isinstance(f, QgsField):
        return f
    return QgsField(f[0], TYPE_MAP.get(f[1], QVariant.String))


def _fieldName(f):
    if isinstance(f, basestring):
        return f
    return f.name()


def newPointsLayer(filename, fields, crs, encoding="utf-8"):
    return newVectorLayer(filename, fields, QGis.WKBPoint, crs, encoding)


def newLinesLayer(filename, fields, crs, encoding="utf-8"):
    return newVectorLayer(filename, fields, QGis.WKBLine, crs, encoding)


def newPolygonsLayer(filename, fields, crs, encoding="utf-8"):
    return newVectorLayer(filename, fields, QGis.WKBPolygon, crs, encoding)


def newVectorLayer(filename, fields, geometryType, crs, encoding="utf-8"):
    if isinstance(crs, basestring):
        crs = QgsCoordinateReferenceSystem(crs)
    if filename is None:
        uri = self.GEOM_TYPE_MAP[geometryType]
        if crs.isValid():
            uri += '?crs=' + crs.authid() + '&'
        fieldsdesc = ['field=' + f for f in fields]

        fieldsstring = '&'.join(fieldsdesc)
        uri += fieldsstring
        layer = QgsVectorLayer(uri, "mem_layer", 'memory')
    else:
        formats = QgsVectorFileWriter.supportedFiltersAndFormats()
        OGRCodes = {}
        for (key, value) in formats.items():
            extension = unicode(key)
            extension = extension[extension.find('*.') + 2:]
            extension = extension[:extension.find(' ')]
            OGRCodes[extension] = value

        extension = os.path.splitext(filename)[1][1:]
        if extension not in OGRCodes:
            extension = 'shp'
            filename = filename + '.shp'

        qgsfields = QgsFields()
        for field in fields:
            qgsfields.append(_toQgsField(field))

        QgsVectorFileWriter(filename, encoding, qgsfields,
                            geometryType, crs, OGRCodes[extension])

        layer = QgsVectorLayer(filename, os.path.basename(filename), 'ogr')

    return VectorLayer(layer)


class Geometry(QgsGeometry):

    def __init__(self, geom):
        self.geom = geom

    def __getattr__(self, name):
        return getattr(self.__dict__['geom'], name)

    def __iter__(self):
        geomType = self.wkbType()
        if geomType in [QGis.WKBMultiPolygon, QGis.WKBMultiPolygon25D]:
            return self.asMultiPolygon()
        elif geomType in [QGis.WKBPolygon, QGis.WKBPolygon25D]:
            return self.asPolygon()
        elif geomType in [QGis.WKBMultiLineString, QGis.WKBMultiLineString25D]:
            return self.asMultiPolyline()
        elif geomType in [QGis.WKBLineString, QGis.WKBLineString25D]:
            return self.asPolyline()
        elif geomType in [QGis.WKBMultiPoint, QGis.WKBMultiPoint25D]:
            return self.asMultiPoint()
        elif geomType in [QGis.WKBPoint, QGis.WKBPoint25D]:
            return self.asPoint()

    def next(self):
        return


def layerFromName(name):
    layers = QgsMapLayerRegistry.instance().mapLayers().values()
    for layer in layers:
        if layer.name() == name:
            return VectorLayer(layer)


def loadLayer(filename):
    name = os.path.split(filename)[1]
    qgslayer = QgsVectorLayer(filename, name, 'ogr')
    if not qgslayer.isValid():
        qgslayer = QgsRasterLayer(filename, name)
        if not qgslayer.isValid():
            raise RuntimeError('Could not load layer: ' + unicode(filename))

    return VectorLayer(qgslayer)