How to access Params from Post-Create Hook?

[s2t2]

Description

I am using this package with SQLAlchemy models. My goal is to use a factory to create a model instance, along with a number of associated model instances. Both specific and generic associations, like these two use cases:

# use two new gyms:
UserFactory(gyms_count=2)

# use specified gym(s)
gym = GymFactory()
UserFactory(gyms=[gym])

I am able to use a post_generation hook to create the related objects, but I'm having issues accessing one of the params when doing so. Or maybe I'm misunderstanding how the params work.

To Reproduce

Model / Factory code

Models (many to many association: user has many gyms, and vice versa):

class Gym(db.Model):
    __tablename__ = "gyms"

    id = db.Column(db.Integer, primary_key=True, index=True)
    title = db.Column(db.String, nullable=False) #> "My Gym"

    memberships = db.relationship("Membership", back_populates="gym")

class User(db.Model):
    __tablename__ = "users"

    id = db.Column(db.Integer, primary_key=True, index=True)
    email = db.Column(db.String, index=True, nullable=False, unique=True)

    memberships = db.relationship("Membership", back_populates="user")

class Membership(db.Model, MyModel):
    __tablename__ = "memberships"

    id = db.Column(db.Integer, primary_key=True, index=True)
    gym_id = db.Column(db.Integer, db.ForeignKey("gyms.id"), nullable=False, index=True)
    user_id = db.Column(db.Integer, db.ForeignKey("users.id"), nullable=False, index=True)

    gym = db.relationship("Gym", back_populates="memberships", uselist=False)
    user = db.relationship("User", back_populates="memberships", uselist=False)

Factories:

class BaseFactory(SQLAlchemyModelFactory):
    class Meta(object):
        sqlalchemy_session = db.session

class UserFactory(BaseFactory):
    class Meta:
        model = User

    id = Sequence(lambda n: n+1)
    email = Sequence(lambda n: f"u{n+1}@example.com")

    class Params:
        gyms_count = 0

    @post_generation
    def gyms(obj, create, extracted, **kwargs):
        if not create:
            # Simple build, do nothing.
            return

        if extracted:
            for gym in extracted:
                Membership.find_or_create(gym_id=gym.id, user_id=obj.id)

        gyms_count = kwargs.get("gyms_count") or 0
        #gyms_count = obj.gyms_count
        # I tried `kwargs.get("gyms_count")` and `obj.gyms_count` but neither was successful.
        # how to get the gyms count param here?

        for _ in range(0, gyms_count):
            gym = GymFactory()
            Membership.find_or_create(gym_id=gym.id, user_id=obj.id)

The issue

How to access the gyms_count param from within the post_generation hook?

[LucasCoderT]

I came across needing this exact kind of situation.

What I've done was just to create my own PostGeneraion subclass:

class PatchedPostGeneration(PostGeneration):
    """
    We need This version of the class because the base class does not pass step to the function

    Step is needed to allow more fined grained control over the post creation of the object.

    """

    def evaluate(self, instance, step, extra):
        """

        Needed by parent class but don't need to implement for our use case

        """
        pass

    def call(self, instance, step, context):
        create = step.builder.strategy == enums.CREATE_STRATEGY
        return self.function(instance, step, create, context)

Then its just used with:

myvalue = PatchedPostGeneration(my_func)

This basically just modifies the default one to pass in an additional step parameter. From the step you can get Params like: step.attributes.get("gyms_count")

As a potential PR idea however is to maybe add it as part of the extra dict. But this is just what I wipped up quickly for my needs.

[n1ngu]

You could always use deep context on the gyms declaration like

class UserFactory(BaseFactory):
    class Meta:
        model = User

    id = Sequence(lambda n: n+1)
    email = Sequence(lambda n: f"u{n+1}@example.com")

    @post_generation
    def gyms(obj, create, extracted, **kwargs):
        if not create:
            # Simple build, do nothing.
            return

        if extracted:
            for gym in extracted:
                Membership.find_or_create(gym_id=gym.id, user_id=obj.id)

        gyms_count = kwargs.pop("count", 0)

        for _ in range(0, gyms_count):
            gym = GymFactory()
            Membership.find_or_create(gym_id=gym.id, user_id=obj.id)

# Note the double underscore so `count` is deep context for the `gyms` declaration
# instead of a gyms_count declaration itself
UserFactory.create(gyms__count=2)

[levic]

I believe this is a duplicate of #544