Fix Pandas datetime and timedelta handling (#48)

* Change pandas handling to use string formatting for datetime. Add corresponding test.

* Fixing timedelta handling, and corresponding test

Author: trangevi

inference_schema/parameter_types/pandas_parameter_type.py

@@ -73,6 +73,9 @@ def deserialize_input(self, input_data):
         if self.enforce_column_type:
             sample_input_column_types = self.sample_input.dtypes.to_dict()
             converted_types = {x: sample_input_column_types.get(x, object) for x in data_frame.columns}
+            for column_name, column_type in converted_types.items():
+                if str(column_type).startswith('timedelta'):
+                    data_frame[column_name] = pd.to_timedelta(data_frame[column_name])
             data_frame = data_frame.astype(dtype=converted_types)
 
         if self.enforce_shape:
@@ -101,11 +104,44 @@ def input_to_swagger(self):
         :rtype: dict
         """
         LIST_LIKE_ORIENTS = ('records', 'values')
-        json_sample = json.loads(self.sample_input.to_json(orient=self.orient))
+        json_sample = json.loads(self.sample_input.to_json(orient=self.orient, date_format='iso'))
 
         if self.orient in LIST_LIKE_ORIENTS:
             swagger_schema = get_swagger_for_list(json_sample)
         else:
             swagger_schema = get_swagger_for_nested_dict(json_sample)
 
+        if self.orient == 'records':
+            for column_name in self.sample_input.columns:
+                data_type = str(self.sample_input.dtypes[column_name])
+                if data_type.startswith('datetime'):
+                    swagger_schema['items']['properties'][str(column_name)]['format'] = 'date-time'
+                elif data_type.startswith('timedelta'):
+                    swagger_schema['items']['properties'][str(column_name)]['format'] = 'timedelta'
+        elif self.orient == 'index':
+            for row in swagger_schema['properties'].values():
+                for column_name in self.sample_input.columns:
+                    data_type = str(self.sample_input.dtypes[column_name])
+                    if data_type.startswith('datetime'):
+                        row['properties'][str(column_name)]['format'] = 'date-time'
+                    elif data_type.startswith('timedelta'):
+                        row['properties'][str(column_name)]['format'] = 'timedelta'
+        elif self.orient == 'columns':
+            for column_name in self.sample_input.columns:
+                for row_info in swagger_schema['properties'][str(column_name)]['properties'].values():
+                    data_type = str(self.sample_input.dtypes[column_name])
+                    if data_type.startswith('datetime'):
+                        row_info['format'] = 'date-time'
+                    elif data_type.startswith('timedelta'):
+                        row_info['format'] = 'timedelta'
+        elif self.orient == 'table':
+            for column_name in self.sample_input.columns:
+                data_type = str(self.sample_input.dtypes[column_name])
+                if data_type.startswith('datetime'):
+                    swagger_schema['properties']['data']['items']['properties'][str(column_name)]['format'] = \
+                        'date-time'
+                elif data_type.startswith('timedelta'):
+                    swagger_schema['properties']['data']['items']['properties'][str(column_name)]['format'] = \
+                        'timedelta'
+
         return swagger_schema

tests/conftest.py

@@ -84,7 +84,8 @@ def pandas_func(param):
 @pytest.fixture(scope="session")
 def decorated_pandas_datetime_func():
     pandas_sample_timestamp_input = pd.DataFrame({'datetime': pd.Series(['2013-12-31T00:00:00.000Z'],
-                                                                        dtype='datetime64[ns]')})
+                                                                        dtype='datetime64[ns]'),
+                                                  'days': pd.Series([pd.Timedelta(days=1)])})
 
     @input_schema('param', PandasParameterType(pandas_sample_timestamp_input))
     def pandas_datetime_func(param):

tests/resources/sample_pandas_datetime_schema.json

@@ -0,0 +1,33 @@
+{
+    "type": "object",
+	"properties": {
+		"param": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "required": [
+                    "datetime",
+                    "days"
+                ],
+                "properties": {
+                    "datetime": {
+                        "type": "string",
+                        "format": "date-time"
+                    },
+                    "days": {
+                        "type": "string",
+                        "format": "timedelta"
+                    }
+                }
+            }
+        }
+    },
+    "example": {
+        "param": [
+            {
+                "datetime": "2013-12-31T00:00:00.000Z",
+                "days": "P1DT0H0M0S"
+            }
+        ]
+    }
+}

tests/test_pandas_parameter_type.py

@@ -33,7 +33,8 @@ def test_pandas_orient_handling(self, decorated_pandas_func_split_orient):
 
     def test_pandas_timestamp_handling(self, decorated_pandas_datetime_func):
         datetime_str = '2013-12-31 00:00:00,000000'
-        pandas_input = {'param': [{'datetime': datetime_str}]}
+        timedelta_str = 'P1DT0H0M0S'
+        pandas_input = {'param': [{'datetime': datetime_str, 'days': timedelta_str}]}
         datetime = pd.DataFrame(
             pd.DataFrame({'datetime': pd.Series([datetime_str], dtype='datetime64[ns]')})['datetime'])
         result = decorated_pandas_datetime_func(**pandas_input)

tests/test_schema_generation.py

@@ -33,6 +33,14 @@ def test_pandas_handling(self, decorated_pandas_func):
         assert ordered(get_output_schema(decorated_pandas_func)) == ordered(self.pandas_sample_output_schema)
 
 
+class TestPandasDatetimeSchemaGeneration(object):
+    pandas_sample_datetime_schema = json.loads(
+        resource_string(__name__, os.path.join('resources', 'sample_pandas_datetime_schema.json')).decode('ascii'))
+
+    def test_pandas_datetime_handling(self, decorated_pandas_datetime_func):
+        assert ordered(get_input_schema(decorated_pandas_datetime_func)) == ordered(self.pandas_sample_datetime_schema)
+
+
 class TestSparkSchemaGeneration(object):
     spark_sample_input_schema = json.loads(
         resource_string(__name__, os.path.join('resources', 'sample_spark_input_schema.json')).decode('ascii'))