0

I'm curious if I could augment/extend the functionality of a base function using functools techniques.

I'm putting together a module with many methods returning generators, and I'm looking for an effective way to dispatch them (?) and also provide the ability to return list without modifying the underlying code.

For example, I might have a simple function which returns a generator. I'd like to extend this function to also have an argument which allows the user to return a list as well.

The following works:

from typing import (
    Union, 
    Generator, 
    List,
    Literal,
)

from functools import wraps
from multimethod import multimethod

# already exists
def base_func(
    a: int, 
    b: int, 
) -> Generator:
    assert b > a
    for i in range(a, b):
        yield i

#newly created
def extended_func(
    a: int, 
    b: int, 
    return_as_list: bool,
) -> Union[Generator, List]:

    gen = base_func(a, b)

    return list(gen) if return_as_list else gen

This seems messy, or at least I'm wondering if it could be done more effectively using dispatching, since I'd like to avoid creating a completely new function.

I scratched together the code below, which does not work (it returns TypeError: type 'bool' is not an acceptable base type), but might be able to get across what I'm trying to do. I know at this point it might obfuscate the original purpose of the function, but I guess I'm just trying to figure out if it's even possible for my own edification.

## decorator
def as_list(func):
    """
    convert generator to list
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        return list(func(*args, **kwargs))
    return wrapper


@multimethod
def test_func(
    a: int, 
    b: int, 
    return_as_list: Literal[False],
) -> Generator[int, None, None]:
    assert b > a
    for i in range(a, b):
        yield i

@test_func.register
@as_list
def _(
    a: int, 
    b: int, 
    return_as_list: Literal[True],
):
    ...
Le Chase
  • 170
  • 9
  • Forgive me for questioning the premise, but why would you need something so elaborate, just to have the option of returning lists instead of generators? If any iterator `g` is finite, it can _always_ be consumed into a list by simply doing `list(g)`, as you demonstrate in your own code. What benefit do you provide to the user with your additional option? Instead of doing `list(base_func(...))` he would now have to do `base_func(..., return_as_list=True)`. Is that _really_ better? – Daniil Fajnberg Mar 01 '23 at 00:00
  • That is a good point, I guess I was just curious. One use case is I might need functionality to switch the default return type of all my generator funcs (of which there are many) to list. For now I've decided to accomplish this with a class (i.e. `base_func` and all other functions are instance methods with the return type an instance attribute). Nevertheless I do agree that this might be a bit opaque, and allowing the user to do `list(base_func(...))` is much simpler. Thanks for your response – Le Chase Mar 01 '23 at 15:33

0 Answers0