Merge 8529ba3 into 80d4010

mapbox · Jan 31, 2018 · 31384d4 · 31384d4
2 parents 80d4010 + 8529ba3
commit 31384d4
Showing 1 changed file with 41 additions and 33 deletions.
diff --git a/mapboxgl/utils.py b/mapboxgl/utils.py
@@ -7,51 +7,59 @@ def df_to_geojson(df, properties=None, lat='lat', lon='lon', precision=None, fil
     """
     geojson = {'type': 'FeatureCollection', 'features': []}
 
+    if filename:
+        with open(filename, 'w+') as f:
+            # Overwrite file if it already exists
+            pass
+
     if precision:
         df[lat] = df[lat].round(precision)
         df[lon] = df[lon].round(precision)
 
     if not properties:
         properties = [c for c in df.columns if c not in [lat, lon]]
 
-    for _, row in df.iterrows():
-        feature = {
-            'type': 'Feature',
-            'properties': {},
-            'geometry': {
-                'type': 'Point',
-                'coordinates': [row[lon], row[lat]]}
-        }
-
-        # TODO performance
-        for prop in properties:
-            feature['properties'][prop] = row[prop]
-
-        geojson['features'].append(feature)
-
     if filename:
-        with open(filename, 'w+') as f:
-            # Overwrite file if it already exists
-            pass
         with open(filename, 'a+') as f:
-            # Write out dictionary contents to a geojson format file
             f.write('{"type": "FeatureCollection", "features": [\n')
-            if len(geojson['features']) > 0:
-                for idx, feat in enumerate(geojson['features']):
-                    if idx == 0:
-                        f.write(json.dumps(feat, ensure_ascii=False,
-                                           sort_keys=True) + '\n')
-                    else:
-                        f.write(',' + json.dumps(feat,
-                                                 ensure_ascii=False, sort_keys=True) + '\n')
+            rowcount = 0
+            for idx, row in df.iterrows():
+                feature = {
+                    'type': 'Feature',
+                    'properties': {},
+                    'geometry': {
+                        'type': 'Point',
+                        'coordinates': [row[lon], row[lat]]}
+                }
+                for prop in properties:
+                    feature['properties'][prop] = row[prop]
+                if rowcount == 0:
+                    f.write(json.dumps(feature, ensure_ascii=False, sort_keys=True) + '\n')
+                    rowcount+=1
+                else:
+                    f.write(',' + json.dumps(feature,
+                                             ensure_ascii=False, sort_keys=True) + '\n')
+                    rowcount+=1
             f.write(']}')
-
-        return {
-            "type": "file",
-            "filename": filename,
-            "feature_count": len(geojson['features'])
-        }
+            return {
+                "type": "file",
+                "filename": filename,
+                "feature_count": rowcount
+            }
     else:
+        for idx, row in df.iterrows():
+            feature = {
+                'type': 'Feature',
+                'properties': {},
+                'geometry': {
+                    'type': 'Point',
+                    'coordinates': [row[lon], row[lat]]}
+            }
+            for prop in properties:
+                feature['properties'][prop] = row[prop]
+
+            geojson['features'].append(feature)
+
         return geojson