From 0e671667b073ee8cfddc4a2633a14dbf44728802 Mon Sep 17 00:00:00 2001
From: Jacob Wasserman <jwasserman@gmail.com>
Date: Fri, 1 Nov 2013 12:58:47 -0400
Subject: [PATCH] ENH: Introduce method to handle NaNs in to_json()

This adds an 'na' option specifying how NaN values should be handled in
GeoDataFrame's to_json(). The default is 'null' which outputs JSON null.
The other options are 'drop' which removes any missing property, and 'keep'
which will output them as NaN.
---
 geopandas/geodataframe.py  | 40 +++++++++++++++++++++++++++++---------
 tests/test_geodataframe.py | 25 +++++++++++++++++++++---
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py
index 2fc01f3..fc41ff6 100644
--- a/geopandas/geodataframe.py
+++ b/geopandas/geodataframe.py
@@ -107,22 +107,44 @@ class GeoDataFrame(DataFrame):
                      coerce_float, params)
 
 
-    def to_json(self, omitna=False, **kwargs):
+    def to_json(self, na='null', **kwargs):
         """Returns a GeoJSON representation of the GeoDataFrame.
 
         Parameters
         ----------
-        omitna : boolean, default False
-            Indicates whether null properties should be included in the 
-            output. This applies to each feature individually so that
-            some features may have different numbers of features if some
-            have null elements
+        na : {'null', 'drop', 'keep'}, default 'null'
+            Indicates how to output missing (NaN) values in the GeoDataFrame
+            * null: ouput the missing entries as JSON null
+            * drop: remove the property from the feature. This applies to
+                    each feature individually so that features may have
+                    different properties
+            * keep: output the missing entries as NaN
         
-        The *kwargs* are passed to json.dumps().
+        The remaining *kwargs* are passed to json.dumps().
         """
+        def fill_none(row):
+            """
+            Takes in a Series, converts to a dictionary with null values
+            set to None
+
+            """
+            na_keys = row.index[row.isnull()]
+            d = row.to_dict()
+            for k in na_keys:
+                d[k] = None
+            return d
+
+        # na_methods must take in a Series and return dict-like
+        na_methods = {'null': fill_none,
+                      'drop': lambda row: row.dropna(),
+                      'keep': lambda row: row}
+
+        if na not in na_methods:
+            raise ValueError('Unknown na method {}'.format(na))
+        f = na_methods[na]
+
         def feature(i, row):
-            if omitna:
-                row = row.dropna()
+            row = f(row)
             return {
                 'id': str(i),
                 'type': 'Feature',
diff --git a/tests/test_geodataframe.py b/tests/test_geodataframe.py
index df1b76f..ca69057 100644
--- a/tests/test_geodataframe.py
+++ b/tests/test_geodataframe.py
@@ -89,13 +89,13 @@ class TestDataFrame(unittest.TestCase):
             props = f['properties']
             self.assertEqual(len(props), 4)
             if props['BoroName'] == 'Queens':
-                self.assertTrue(np.isnan(props['Shape_Area']))
+                self.assertTrue(props['Shape_Area'] is None)
 
-    def test_to_json_omitna(self):
+    def test_to_json_dropna(self):
         self.df['Shape_Area'][self.df['BoroName']=='Queens'] = np.nan
         self.df['Shape_Leng'][self.df['BoroName']=='Bronx'] = np.nan
 
-        text = self.df.to_json(omitna=True)
+        text = self.df.to_json(na='drop')
         data = json.loads(text)
         self.assertEqual(len(data['features']), 5)
         for f in data['features']:
@@ -113,6 +113,25 @@ class TestDataFrame(unittest.TestCase):
             else:
                 self.assertEqual(len(props), 4)
 
+    def test_to_json_keepna(self):
+        self.df['Shape_Area'][self.df['BoroName']=='Queens'] = np.nan
+        self.df['Shape_Leng'][self.df['BoroName']=='Bronx'] = np.nan
+
+        text = self.df.to_json(na='keep')
+        data = json.loads(text)
+        self.assertEqual(len(data['features']), 5)
+        for f in data['features']:
+            props = f['properties']
+            self.assertEqual(len(props), 4)
+            if props['BoroName'] == 'Queens':
+                self.assertTrue(np.isnan(props['Shape_Area']))
+                # Just make sure setting it to nan in a different row
+                # doesn't affect this one
+                self.assertTrue('Shape_Leng' in props)
+            elif props['BoroName'] == 'Bronx':
+                self.assertTrue(np.isnan(props['Shape_Leng']))
+                self.assertTrue('Shape_Area' in props)
+
     def test_copy(self):
         df2 = self.df.copy()
         self.assertTrue(type(df2) is GeoDataFrame)