Skip to content

Use consumer generator method decorator to attach type annotations to created generator. #171

@cboulay

Description

@cboulay

I'm working on a function to check that a chain of processors all have compatible message types. For new processors, this works fine. It breaks down when one of the processors in the chain is an older processor using a generator. It's possible but unreliable to get the generator's origin method's signature from a globals lookup. A preferable solution would be to attach the annotations directly to the generator when it is created from the generator function when the annotations are easily accessible. This isn't very difficult in ezmsg because we almost always decorate our generator methods. Something like the following might allow the type annotations to be attached to the created generator.

Possible change:

def consumer(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        gen = func(*args, **kwargs)
        # Extract type hints from the generator function
        ret_type = typing.get_type_hints(func).get('return')
        if ret_type and hasattr(ret_type, "__origin__") and ret_type.__origin__ is typing.Generator:
            gen_args = typing.get_args(ret_type)
            if len(gen_args) >= 2:
                # Attach type information directly to the generator instance
                gen._generator_types = (gen_args[0], gen_args[1])
        next(gen)  # Prime the generator
        return gen
    return wrapper

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions