Refactor PlotlyJSONEncoder. (see below)

theengineear · theengineear · commit 0bd4dc307851 · 2014-12-15T09:28:46.000-08:00
Detiails:

- renamed *classy* names to *methody* names
- made methods `static`
- added docstrings to all the things
- removed some replication of code
- modified the `sage` import to be more explicit
- some reordering of the encoders (see TODO)
- let `json.JSONEncoder.default` throw `TypeError`
diff --git a/plotly/utils.py b/plotly/utils.py
@@ -26,7 +26,7 @@
     _pandas_imported = False
 
 try:
-    from sage.all import RR, ZZ
+    import sage.all
     _sage_imported = True
 except ImportError:
     _sage_imported = False
@@ -91,109 +91,147 @@ def ensure_dir_exists(directory):
             os.makedirs(directory)
 
 
+def iso_to_plotly_time_string(iso_string):
+    """Remove timezone info and replace 'T' delimeter with ' ' (ws)."""
+    no_tz_string = '-'.join(iso_string.split('-')[:3])
+    if no_tz_string.endswith('T00:00:00'):
+        return no_tz_string.replace('T00:00:00', '')
+    else:
+        return no_tz_string.replace('T', ' ')
+
+
 ### Custom JSON encoders ###
 class NotEncodable(Exception):
     pass
 
 
-    def numpyJSONEncoder(self, obj):
-        if not _numpy_imported:
-            raise NotEncodable
 class PlotlyJSONEncoder(json.JSONEncoder):
+    """
+    Meant to be passed as the `cls` kwarg to json.dumps(obj, cls=..)
 
-        if obj is numpy.ma.core.masked:
-            return float('nan')
+    See PlotlyJSONEncoder.default for more implementation information.
 
-        if type(obj).__module__.split('.')[0] == numpy.__name__:
-            return obj.tolist()
+    """
+    @staticmethod
+    def encode_as_plotly(obj):
+        """Attempt to use a builtin `ploty_to_json` method."""
+        try:
+            return obj.to_plotly_json()
+        except AttributeError:
+            raise NotEncodable
 
+    @staticmethod
+    def encode_as_list(obj):
+        """Attempt to use `tolist` method to convert to normal Python list."""
+        if hasattr(obj, 'tolist'):
+            return obj.tolist()
         else:
             raise NotEncodable
 
-    def datetimeJSONEncoder(self, obj):
-        """
-        if datetime or iterable of datetimes,
-        convert to a string that plotly understands
-        format as %Y-%m-%d %H:%M:%S.%f,
-                  %Y-%m-%d %H:%M:%S, or
-                  %Y-%m-%d
-        depending on what non-zero resolution was provided
-        """
-
-        if _pandas_imported and obj is pandas.NaT:
-            return None
-
+    @staticmethod
+    def encode_as_sage(obj):
+        """Attempt to convert sage.all.RR to floats and sage.all.ZZ to ints"""
+        if not _sage_imported:
+            raise NotEncodable
 
-        if isinstance(obj, datetime.datetime):
-            if obj.microsecond:
-                return obj.strftime('%Y-%m-%d %H:%M:%S.%f')
-            elif any((obj.second, obj.minute, obj.hour)):
-                return obj.strftime('%Y-%m-%d %H:%M:%S')
-            else:
-                return obj.strftime('%Y-%m-%d')
-        elif isinstance(obj, datetime.date):
-            return obj.strftime('%Y-%m-%d')
+        if obj in sage.all.RR:
+            return float(obj)
+        elif obj in sage.all.ZZ:
+            return int(obj)
         else:
             raise NotEncodable
 
-    def pandasJSONEncoder(self, obj):
+    @staticmethod
+    def encode_as_pandas(obj):
+        """Attempt to convert pandas.NaT"""
         if not _pandas_imported:
             raise NotEncodable
 
-        if isinstance(obj, pandas.Series):
-            serializable_list = []
-            for li in list(obj):
-                try:
-                    json.dumps(li)
-                    serializable_list.append(li)
-                except TypeError:
-                    serializable_list.append(self.default(li))
-
-            return serializable_list
-        elif isinstance(obj, pandas.Index):
-            return obj.tolist()
-        elif obj is pandas.NaT:
+        if obj is pandas.NaT:
             return None
         else:
             raise NotEncodable
 
-    def sageJSONEncoder(self, obj):
-        if not _sage_imported:
+    @staticmethod
+    def encode_as_numpy(obj):
+        """Attempt to convert numpy.ma.core.masked"""
+        if not _numpy_imported:
             raise NotEncodable
 
-        if obj in RR:
-            return float(obj)
-        elif obj in ZZ:
-            return int(obj)
+        if obj is numpy.ma.core.masked:
+            return float('nan')
         else:
             raise NotEncodable
 
-    def builtinJSONEncoder(self, obj):
-        '''
-        Provide an API for folks to write their own
-        JSON encoders for their objects that they wanna
-        send to Plotly: this is an object's `to_plotly_json` method
+    @staticmethod
+    def encode_as_datetime(obj):
+        """Attempt to convert to utc-iso time string using datetime methods."""
+        try:
+            if obj.utcoffset():
+                obj = obj - obj.utcoffset()
+            time_string = obj.isoformat()
+        except AttributeError:
+            raise NotEncodable
+        else:
+            return iso_to_plotly_time_string(time_string)
 
-        Used for grid_objs.Column objects.
-        '''
+    @staticmethod
+    def encode_as_date(obj):
+        """Attempt to convert to utc-iso time string using date methods."""
         try:
-            return obj.to_plotly_json()
+            time_string = obj.isoformat()
         except AttributeError:
             raise NotEncodable
+        else:
+            return iso_to_plotly_time_string(time_string)
 
     def default(self, obj):
-        encoders = (self.builtinJSONEncoder,
-                    self.datetimeJSONEncoder,
-                    self.numpyJSONEncoder,
-                    self.pandasJSONEncoder,
-                    self.sageJSONEncoder)
-        for encoder in encoders:
+        """
+        Accept an object (of unknown type) and try to encode with priority:
+        1. builtin:     user-defined objects
+        2. sage:        sage math cloud
+        3. pandas:      dataframes/series
+        4. numpy:       ndarrays
+        5. datetime:    time/datetime objects
+
+        Each method throws a NotEncoded exception if it fails.
+
+        The default method will only get hit if the object is not a type that
+        is naturally encoded by json:
+
+            Normal objects:
+                dict                object
+                list, tuple         array
+                str, unicode        string
+                int, long, float    number
+                True                true
+                False               false
+                None                null
+
+            Extended objects:
+                float('nan')        'NaN'
+                float('infinity')   'Infinity'
+                float('-infinity')  '-Infinity'
+
+        Therefore, we only anticipate either unknown iterables or values here.
+
+        """
+        # TODO: The ordering if these methods is *very* important. Is this OK?
+        encoding_methods = (
+            self.encode_as_plotly,
+            self.encode_as_sage,
+            self.encode_as_numpy,
+            self.encode_as_pandas,
+            self.encode_as_datetime,
+            self.encode_as_date,
+            self.encode_as_list  # because some values have `tolist` do last.
+        )
+        for encoding_method in encoding_methods:
             try:
-                return encoder(obj)
+                return encoding_method(obj)
             except NotEncodable:
                 pass
-
-        raise TypeError(repr(obj) + " is not JSON serializable")
+        return json.JSONEncoder.default(self, obj)
 
 
 ### unicode stuff ###
