DRF: DataFrame を使って CSV エクスポート

CSV レンダラー

1
2
3
4
5
6
7
def dataframe_to_csv_stream(df, encoding=None, **kwargs):
    buffer = BytesIO()
    df.to_csv(buffer, **kwargs)
    buffer.seek(0)
    if encoding:
        return BytesIO(buffer.read().decode("utf8").encode(encoding))
    return buffer
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class DataFrameCSVRenderer(BaseRenderer):
    media_type = "text/csv"
    format = "csv"

    def render(self, data, media_type=None, renderer_context=None, writer_opts=None):
        renderer_context = renderer_context or {}
        if not isinstance(data, pd.DataFrame):
            data = pd.DataFrame(data)
        stream = dataframe_to_csv_stream(data, index=False)
        return stream.read().decode("utf-8-sig")

DataFrame を返すリストシリアライザ

1
2
3
4
5
6
7
class DataFrameSerializer(serializers.ListSerializer):
    @property
    def data(self):
        ret = super().data
        ser = ret.serializer.child
        columns = dict((name, field.label) for name, field in ser.fields.items())
        return pd.DataFrame(ret).rename(columns=columns)

エクスポートするシリアライザを生成するメタクラス

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

class ExpporSerializerMetaclass(serializers.SerializerMetaclass):
    EXPORT_META = {}
    FIELD_CLASS_DEFAULTS = {}

    @classmethod
    def create_fields(cls):
        def _create_field(item):
            label, name, klass = item
            defaults = cls.FIELD_CLASS_DEFAULTS[klass]
            return (name, getattr(serializers, klass)(label=label, **defaults))

        return dict(map(_create_field, cls.EXPORT_META))

    @classmethod
    def annotate(cls, queryset):
        raise NotImplementedError()

    def __new__(cls, name, bases, attrs, **kwargs):
        attrs["Meta"].fields = [i[1] for i in cls.EXPORT_META]      # CSVで返答するカラム
        attrs["Meta"].list_serializer_class = DataFrameSerializer   # リストシリアライザ
        attrs.update(**cls.create_fields())
        attrs["annotate"] = cls.annotate    # クエリセットをアノテートさせるクラスメソッド 

        return super().__new__(cls, name, bases, attrs, **kwargs)

実際のメタクラス

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class IncomeExpporSerializerMetaclass(ExpporSerializerMetaclass):
    # (ラベル, フィールド名, フィールドクラス名)
    EXPORT_META = [
        ("会社コード", "order__company_code", "CharField"),
        ....
    ]

    # フィールドクラスのデフォルト値
    FIELD_CLASS_DEFAULTS = dict(
        CharField=dict(read_only=True),
        DecimalField=dict(max_digits=12, decimal_places=1, read_only=True),
    )

    @classmethod
    def annotate(cls, queryset):
        """ クエリセットを適切にアノテートする"""
        def _ann(item):
            label, name, klass = item
            if name.startswith("order__"):
                # リレーション
                return (name, F(name))
            if name.startswith("meme__"):
                # 固定値
                return (name, Value(""))
            # その他はモデルフィールドなのでアノテートしない
            return (None, None)

        ann = dict(filter(lambda i: i[1] is not None, map(_ann, cls.EXPORT_META)))
        return queryset.annotate(**ann)

実際のシリアライザ

1
2
3
class IncomeExportSerializer(serializers.ModelSerializer, metaclass=IncomeExpporSerializerMetaclass):
    class Meta:
        model = models.Income

API 本体(ビューセット)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class IncomeViewSet(VS.BaseModelViewSet):

    serializer_class = serializers.IncomeSerializer
    queryset = models.Income.objects.all()
    pagination_class = paginations.Pagination

    renderer_classes = tuple(api_settings.DEFAULT_RENDERER_CLASSES) + (DataFrameCSVRenderer,)

    def get_serializer_class(self):
        """ export アクションでシリアライザを切り替える """
        res = {
            "export": serializers.IncomeExportSerializer,
        }.get(self.action, super().get_serializer_class())
        return res

    def get_queryset(self):
        """ export アクションの場合、必要なアノテーションを行う """
        qs = super().get_queryset()
        if self.action == "export":
            qs = serializers.IncomeExportSerializer.annotate(qs).filter(level=0)
            return qs
        return qs

    @decorators.action(methods=["get"], detail=False)
    def export(self, request, pk=None):
        ## リスト処理を行うと シリアライザとレンダラで勝手にCSVを返答する(Accept: text/csvの時)
        return self.list(request)