1# -*- coding: utf-8 -*-
2
3"""
4***************************************************************************
5    KeepNBiggestParts.py
6    ---------------------
7    Date                 : July 2014
8    Copyright            : (C) 2014 by Victor Olaya
9    Email                : volayaf at gmail dot com
10***************************************************************************
11*                                                                         *
12*   This program is free software; you can redistribute it and/or modify  *
13*   it under the terms of the GNU General Public License as published by  *
14*   the Free Software Foundation; either version 2 of the License, or     *
15*   (at your option) any later version.                                   *
16*                                                                         *
17***************************************************************************
18"""
19
20__author__ = 'Victor Olaya'
21__date__ = 'July 2014'
22__copyright__ = '(C) 2014, Victor Olaya'
23
24from operator import itemgetter
25
26from qgis.core import (QgsGeometry,
27                       QgsFeatureSink,
28                       QgsProcessing,
29                       QgsProcessingException,
30                       QgsProcessingParameterFeatureSink,
31                       QgsProcessingParameterFeatureSource,
32                       QgsProcessingParameterNumber,
33                       )
34
35from processing.algs.qgis.QgisAlgorithm import QgisAlgorithm
36
37
38class KeepNBiggestParts(QgisAlgorithm):
39    POLYGONS = 'POLYGONS'
40    PARTS = 'PARTS'
41    OUTPUT = 'OUTPUT'
42
43    def group(self):
44        return self.tr('Vector geometry')
45
46    def groupId(self):
47        return 'vectorgeometry'
48
49    def __init__(self):
50        super().__init__()
51
52    def initAlgorithm(self, config=None):
53        self.addParameter(QgsProcessingParameterFeatureSource(self.POLYGONS,
54                                                              self.tr('Polygons'), [QgsProcessing.TypeVectorPolygon]))
55        self.addParameter(QgsProcessingParameterNumber(self.PARTS,
56                                                       self.tr('Parts to keep'),
57                                                       QgsProcessingParameterNumber.Integer,
58                                                       1, False, 1))
59        self.addParameter(
60            QgsProcessingParameterFeatureSink(self.OUTPUT, self.tr('Parts'), QgsProcessing.TypeVectorPolygon))
61
62    def name(self):
63        return 'keepnbiggestparts'
64
65    def displayName(self):
66        return self.tr('Keep N biggest parts')
67
68    def processAlgorithm(self, parameters, context, feedback):
69        source = self.parameterAsSource(parameters, self.POLYGONS, context)
70        if source is None:
71            raise QgsProcessingException(self.invalidSourceError(parameters, self.POLYGONS))
72
73        parts = self.parameterAsInt(parameters, self.PARTS, context)
74
75        fields = source.fields()
76        (sink, dest_id) = self.parameterAsSink(parameters, self.OUTPUT, context,
77                                               source.fields(), source.wkbType(), source.sourceCrs())
78        if sink is None:
79            raise QgsProcessingException(self.invalidSinkError(parameters, self.OUTPUT))
80
81        features = source.getFeatures()
82        total = 100.0 / source.featureCount() if source.featureCount() else 0
83        for current, feat in enumerate(features):
84            if feedback.isCanceled():
85                break
86
87            geom = feat.geometry()
88            if geom.isMultipart():
89                out_feature = feat
90                geoms = geom.asGeometryCollection()
91                geom_area = [(i, geoms[i].area()) for i in range(len(geoms))]
92                geom_area.sort(key=itemgetter(1))
93                if parts == 1:
94                    out_feature.setGeometry(geoms[geom_area[-1][0]])
95                elif parts > len(geoms):
96                    out_feature.setGeometry(geom)
97                else:
98                    out_feature.setGeometry(geom)
99                    geomres = [geoms[i].asPolygon() for i, a in geom_area[-1 * parts:]]
100                    out_feature.setGeometry(QgsGeometry.fromMultiPolygonXY(geomres))
101                sink.addFeature(out_feature, QgsFeatureSink.FastInsert)
102            else:
103                sink.addFeature(feat, QgsFeatureSink.FastInsert)
104
105            feedback.setProgress(int(current * total))
106
107        return {self.OUTPUT: dest_id}
108