[tips] factoryboy で作成したモデルに type hint をつける方法

factoryboy は python 用のモックアップデータを作成するライブラリです。 Django の model object にも対応していて required な field を動的に生成したり (ex. user_0001 みたいに連番にしたり, アルファベットをランダムに選んだ文字列にしたり...)、relation を持つ object を relation 先も一緒に作成できて便利です。

例えば以下のような user / article model があったとしましょう。

from django.db import models


class User(models.Model):
    name = models.CharField(max_length=8, unique=True)


class Article(models.Model):
    author = models.ForeignKey(User, on_delete=models.CASCADE)
    title = models.CharField(max_length=100)
    body = models.TextField(null=True, blank=True)

この時適当な article object を作りたいとなると

作成者になる user を作成して保存.
ブログのタイトルを適当に生成
1,2 をつかって article を作成

という3ステップを踏む必要があります。

user = User(name='foo')
user.save()
article = Article(author=user, title='bar')

これを見ると例えば

いちいち article を作るのに user を作成するのは面倒なので generate_article を作りたくなる
article の初期値は自分で弄りたいので, 引数に user / title / body は含めたい

と思うことでしょう。関数にするならば以下のような形でしょうか。

def generate_article(user = None, title = None, body = None):
    if user is None:
        user = User(name='foo')
        user.save()

    if title is None:
        title = 'bar'

    return Article.objects.create(user=user, title=title, body=body)

さあこれでOK! と思うかもしれませんが、このコードには問題があります。そう User の name は unique 制約が付いているのです。したがって user=None で2回 generate_article を呼び出すとエラーになってしまいます。困りますね。 User を作成する時 name を何らかのランダムな文字列にする、というようなコードを書く必要がありますが、それらをいちいち書いていると大変です。

こういうモックデータを作成するときに使えるのが factoryboy です。factorybody で article を作成するコードを実装すると以下のようになります。

from factory.django import DjangoModelFactory
import factory
from factory.fuzzy import FuzzyText


class UserFactory(DjangoModelFactory):
    class Meta:
        model = User
        django_get_or_create = ("name",)

    name = factory.Sequence(lambda n: f"u-{n}")

class ArticleFacotry(DjangoModelFactory):
    class Meta:
        model = Article
    
    author = factory.SubFactory(UserFactory)
    title = FuzzyText(length=12)

簡単に説明をすると Sequence では作成のたびに関数が呼び出されて引数の n が増えていきます。また FuzzyText ではランダムに文字列が, SubFactory は relation を持つモデルのときに、指定された Factory (上記の場合だと UserFactory) で作成されたモデルが設定されるようになります。このように、作成したタイミングで動的に属性指定をしてモデルを作成する手伝いをしてくれるのが factoryboy です。

ここから article を作るのは以下で終わりです。

article = ArticleFacotry()

簡単ですね。また author だけ指定したい! という時には __init__ に author を渡して上げればOKです。簡単ですね。

goto = UserFactory(name="goto")
goto_article = ArticleFacotry(author=goto)

assert goto_article.author == goto # True

このように便利なのですが1点困るのは pycharm のような IDE で factoryboy の constructor で作成された object が作成対象の object としてみなされないことです。なので普通に作成してインテリセンスを利かせてもほしい物が出てきません。

f:id:dette:20210702004733p:plain — factoryboy object だと思い込んでいる(それはそう)なので、サジェストが django model のそれになっていない。

関連する issue. みんな困ってるみたいですね。
- Support type hinting? · Issue #468 · FactoryBoy/factory_boy · GitHub
- How to add hint to factory method? · Issue #58 · python/typing · GitHub

これを解決する方法の一つとして紹介されていたのが「TypeVar T を定義して, T を返すような class method を用意する」方法です。(ref: https://github.com/FactoryBoy/factory_boy/issues/468#issuecomment-759452373 )

from typing import Generic
from typing import TypeVar
import factory

T = TypeVar("T")

class BaseFactory(Generic[T], factory.django.DjangoModelFactory):
    @classmethod
    def create(cls, **kwargs) -> T:
        return super().create(**kwargs)

使う方ではこんな感じ。

class ArticleFactory(BaseFactory[Article]):
    class Meta:
        model = Article

    author = factory.SubFactory(UserFactory)
    title = FuzzyText(length=12)

元の実装と違うのは BaseFactory に [] で返り値のクラス T を指定している点です。こうすることで T に与えられた object が create の返り値とみなされるため IDE で返り値は django model の instance だと思ってくれます。

# __init__ (constructor) は使わずに
# article = ArticleFacotry() 

# `create` を使う
article = ArticleFactory.create()

f:id:dette:20210702004633p:plain — ちゃんと理解してもらえて嬉しい

author のほうもこのとおり。

というわけで factoryboy の tips でした。

nykergoto’s blog

機械学習とpythonをメインに

[tips] factoryboy で作成したモデルに type hint をつける方法