include pk in default JSON; detect custom pk name

sheppard · sheppard · commit 1955c6d68005 · 2017-09-13T10:50:13.000-05:00
(fixes #29)
diff --git a/rest_pandas/renderers.py b/rest_pandas/renderers.py
@@ -36,16 +36,21 @@ def render(self, data, accepted_media_type=None, renderer_context=None):
             raise Exception(
                 RESPONSE_ERROR % type(data).__name__
             )
+
         name = getattr(self, 'function', "to_%s" % self.format)
-        function = getattr(data, name, None)
-        if not function:
+        if not hasattr(data, name):
             raise Exception("Data frame is missing %s property!" % name)
+
         self.init_output()
         args = self.get_pandas_args(data)
         kwargs = self.get_pandas_kwargs(data, renderer_context)
-        function(*args, **kwargs)
+        self.render_dataframe(data, name, *args, **kwargs)
         return self.get_output()
 
+    def render_dataframe(self, data, name, *args, **kwargs):
+        function = getattr(data, name)
+        function(*args, **kwargs)
+
     def init_output(self):
         self.output = StringIO()
 
@@ -150,19 +155,44 @@ class PandasJSONRenderer(PandasBaseRenderer):
     media_type = "application/json"
     format = "json"
 
+    orient_choices = {
+        'records-index',  # Unique to DRP
+        'split',
+        'records',
+        'index',
+        'columns',
+        'values',
+        'table',
+    }
+    default_orient = 'records-index'
+
+    date_format_choices = {'epoch', 'iso'}
+    default_date_format = 'iso'
+
     def get_pandas_kwargs(self, data, renderer_context):
         request = renderer_context['request']
+
         orient = request.GET.get('orient', '')
+        if orient not in self.orient_choices:
+            orient = self.default_orient
+
         date_format = request.GET.get('date_format', '')
-        if orient not in {'split', 'records', 'index', 'columns', 'values'}:
-            orient = 'records'
-        if date_format not in {'epoch', 'iso'}:
-            date_format = 'iso'
+        if date_format not in self.date_format_choices:
+            date_format = self.default_date_format
+
         return {
             'orient': orient,
             'date_format': date_format,
         }
 
+    def render_dataframe(self, data, name, *args, **kwargs):
+        if kwargs.get('orient') == 'records-index':
+            kwargs['orient'] = 'records'
+            data.reset_index(inplace=True)
+        return super(PandasJSONRenderer, self).render_dataframe(
+            data, name, *args, **kwargs
+        )
+
 
 class PandasExcelRenderer(PandasFileRenderer):
     """
diff --git a/rest_pandas/serializers.py b/rest_pandas/serializers.py
@@ -68,11 +68,17 @@ def get_index_fields(self):
         """
         List of fields to use for index
         """
-        default_fields = []
-        if getattr(self.model_serializer_meta, 'model', None):
-            if 'id' in self.child.get_fields():
-                default_fields = ['id']
-        return self.get_meta_option('index', default_fields)
+        index_fields = self.get_meta_option('index', [])
+        if index_fields:
+            return index_fields
+
+        model = getattr(self.model_serializer_meta, 'model', None)
+        if model:
+            pk_name = model._meta.pk.name
+            if pk_name in self.child.get_fields():
+                return [pk_name]
+
+        return []
 
     def get_meta_option(self, name, default=None):
         meta_name = 'pandas_' + name
diff --git a/tests/test_views.py b/tests/test_views.py
@@ -1,5 +1,5 @@
 from rest_framework.test import APITestCase
-from tests.testapp.models import TimeSeries
+from tests.testapp.models import TimeSeries, CustomIndexSeries
 from wq.io import load_string
 import json
 import datetime
@@ -17,6 +17,10 @@ def setUp(self):
         )
         for date, value in data:
             TimeSeries.objects.create(date=date, value=value)
+            CustomIndexSeries.objects.create(
+                code='v' + date.replace('-', ''),
+                value=value,
+            )
 
     def test_view_csv(self):
         response = self.client.get("/timeseries.csv")
@@ -41,6 +45,8 @@ def test_view_json(self):
         self.assertEqual(response.accepted_media_type, "application/json")
         data = json.loads(response.content.decode('utf-8'))
         self.assertEqual(len(data), 5)
+        self.assertIn('id', data[0])
+        self.assertEqual(data[0]["id"], 1)
         self.assertEqual(data[0]["value"], 0.5)
         self.assertEqual(data[0]["date"], "2014-01-01T00:00:00.000Z")
 
@@ -53,6 +59,11 @@ def test_view_json_kwargs(self):
         date = datetime.datetime.utcfromtimestamp(data[0]["date"] / 1000)
         self.assertEqual(date, datetime.datetime(2014, 1, 1))
 
+        response = self.client.get("/timeseries.json?orient=index")
+        data = json.loads(response.content.decode('utf-8'))
+        self.assertEqual(len(data.values()), 5)
+        self.assertEqual(data["1"]["value"], 0.5)
+
     def test_view_html(self):
         response = self.client.get("/timeseries?test=1")
         expected = open(
@@ -111,5 +122,24 @@ def test_from_file(self):
         self.assertEqual(len(data), 4)
         self.assertEqual(data[0].x, '5')
 
+    def test_customindex_csv(self):
+        response = self.client.get("/customindex.csv")
+        data = self.load_string(response)
+        self.assertEqual(len(data), 5)
+        self.assertEqual(data[0].code, 'v20140101')
+        self.assertEqual(data[0].value, '0.5')
+
+    def test_customindex_json(self):
+        response = self.client.get("/customindex.json")
+        data = json.loads(response.content.decode('utf-8'))
+        self.assertEqual(len(data), 5)
+        self.assertEqual(data[0]['code'], 'v20140101')
+        self.assertEqual(data[0]['value'], 0.5)
+
+        response = self.client.get("/customindex.json?orient=index")
+        data = json.loads(response.content.decode('utf-8'))
+        self.assertEqual(len(data), 5)
+        self.assertEqual(data['v20140101']['value'], 0.5)
+
     def load_string(self, response):
         return load_string(response.content.decode('utf-8'))
diff --git a/tests/testapp/models.py b/tests/testapp/models.py
@@ -33,3 +33,8 @@ class ComplexTimeSeries(models.Model):
     # Values
     value = models.FloatField()
     flag = models.CharField(max_length=1, null=True, blank=True)
+
+
+class CustomIndexSeries(models.Model):
+    code = models.SlugField(primary_key=True)
+    value = models.FloatField()
diff --git a/tests/testapp/serializers.py b/tests/testapp/serializers.py
@@ -1,7 +1,9 @@
 from rest_framework.serializers import ModelSerializer
 from rest_framework import serializers
 from rest_pandas import PandasUnstackedSerializer
-from .models import TimeSeries, MultiTimeSeries, ComplexTimeSeries
+from .models import (
+    TimeSeries, MultiTimeSeries, ComplexTimeSeries, CustomIndexSeries,
+)
 
 
 class TimeSeriesSerializer(ModelSerializer):
@@ -62,3 +64,9 @@ class Meta:
         list_serializer_class = PandasUnstackedSerializer
         # pandas_unstacked_header = Missing
         pandas_index = ['series']
+
+
+class CustomIndexSeriesSerializer(ModelSerializer):
+    class Meta:
+        model = CustomIndexSeries
+        fields = '__all__'
diff --git a/tests/testapp/urls.py b/tests/testapp/urls.py
@@ -8,6 +8,7 @@
     DjangoPandasView, TimeSeriesViewSet,
     MultiTimeSeriesView, MultiScatterView, MultiBoxplotView,
     ComplexTimeSeriesView, ComplexScatterView, ComplexBoxplotView,
+    CustomIndexSeriesView,
 )
 
 router = DefaultRouter()
@@ -28,6 +29,7 @@
     url(r'^complextimeseries$', ComplexTimeSeriesView.as_view()),
     url(r'^complexscatter$', ComplexScatterView.as_view()),
     url(r'^complexboxplot$', ComplexBoxplotView.as_view()),
+    url(r'^customindex$', CustomIndexSeriesView.as_view()),
 ]
 urlpatterns = format_suffix_patterns(urlpatterns)
 urlpatterns += [
diff --git a/tests/testapp/views.py b/tests/testapp/views.py
@@ -5,12 +5,15 @@
 from rest_framework import renderers
 from rest_framework.generics import ListAPIView
 from rest_pandas import renderers as pandas_renderers
-from .models import TimeSeries, MultiTimeSeries, ComplexTimeSeries
+from .models import (
+    TimeSeries, MultiTimeSeries, ComplexTimeSeries, CustomIndexSeries,
+)
 from .serializers import (
     TimeSeriesSerializer, TimeSeriesNoIdSerializer,
     MultiTimeSeriesSerializer,
     ComplexTimeSeriesSerializer, ComplexScatterSerializer,
     ComplexBoxplotSerializer,
+    CustomIndexSeriesSerializer,
 )
 import pandas as pd
 
@@ -122,3 +125,8 @@ class ComplexBoxplotView(PandasView):
     queryset = ComplexTimeSeries.objects.all()
     serializer_class = ComplexBoxplotSerializer
     pandas_serializer_class = PandasBoxplotSerializer
+
+
+class CustomIndexSeriesView(PandasView):
+    queryset = CustomIndexSeries.objects.all()
+    serializer_class = CustomIndexSeriesSerializer

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@`
`8`	`8`	`DjangoPandasView, TimeSeriesViewSet,`
`9`	`9`	`MultiTimeSeriesView, MultiScatterView, MultiBoxplotView,`
`10`	`10`	`ComplexTimeSeriesView, ComplexScatterView, ComplexBoxplotView,`
	`11`	`+ CustomIndexSeriesView,`
`11`	`12`	`)`
`12`	`13`
`13`	`14`	`router = DefaultRouter()`
`@@ -28,6 +29,7 @@`
`28`	`29`	`url(r'^complextimeseries$', ComplexTimeSeriesView.as_view()),`
`29`	`30`	`url(r'^complexscatter$', ComplexScatterView.as_view()),`
`30`	`31`	`url(r'^complexboxplot$', ComplexBoxplotView.as_view()),`
	`32`	`+ url(r'^customindex$', CustomIndexSeriesView.as_view()),`
`31`	`33`	`]`
`32`	`34`	`urlpatterns = format_suffix_patterns(urlpatterns)`
`33`	`35`	`urlpatterns += [`