Add autocomplete support for builds(<target>, populate_full_signature=True)

[rsokl]

This PR:

1. Leverages ParamSpec to enable IDE autocompletion on the configs returned by builds(<target>, populate_full_signature=True).
2. Adds overloads to make_custom_builds_fn to reflect relevant updated default values on the resulting builds function. Specifically, this holds for zen_partial=True and populate_full_signature=True, respectively.
3. Makes ZenWrappers a "public" type via hydra_zen.typing

Adding autocomplete for builds(<target>, populate_full_signature=True)

Consider the following builds pattern:

def f(x: int, y: str = "hello") -> bool:
    return False


Conf = builds(f, populate_full_signature=True)

Presently, IDEs / type-checkers cannot infer the signature to Conf, even though the user can infer that the full signature of f is reflected in Conf. This PR fixes that!

Before:

After:

This means that instances of Conf are now amenable to static type-checking:

Lastly, type-checkers are still able to "see" the correct type returned via instantiate:

Adding overloads to make_custom_builds_fn

Previously, we did not provide type-checkers with "dynamic" information about the builds-function produced by make_custom_builds_fn. Now, type-checkers can "see" some of the new defaults in the builds function created by make_custom_builds_fn.

Before:

partial_builds = make_custom_builds_fn(zen_partial=True)

# type-checker doesn't "see" that `zen_partial=True`
Conf = partial_builds(int)  # type: Type[Builds[Type[int]]]

After

partial_builds = make_custom_builds_fn(zen_partial=True)

# `zen_partial=True` is inferred by type-checker!
Conf = partial_builds(int)  # type: Type[PartialBuilds[Type[int]]]


# `Manually-overriding zen_partial behaves as-expected
Conf = partial_builds(int, zen_partial=False)  # type: Type[Builds[Type[int]]]

This is especially nice, because now the aforementioned improvements to static inference on builds(<target>, populate_full_signature=True) can manifest via make_custom_builds_fn (otherwise, users would always have to write
builds(<target>, populate_full_signature=True) to get the nice static inference):

def f(x: int): return

full_builds = make_custom_builds_fn(populate_full_signature=True)

Conf = full_builds(f)  # type: Conf: (x: int) -> Builds[Type[None]]

Conf(x="hi")  # type-checker: error

This was accomplished by adding special overloads for builds with updated default values -- specifically for populate_full_signature=True and zen_partial=True, respectively -- which are stored in hydra_zen.typing._builds_overloads. We then use some fanciness to conditionally dispatch the output of make_custom_builds_fn using these specialized overloads.

This adds some unfortunate boilerplate and redundancies, but I think it is ultimately well worth it. Otherwise the very common use case of make_custom_builds_fn(populate_full_signature=True) would not be as powerful as manually writing builds(<t>, populate_full_signature=True); furthermore, it would be highly non-trivial for end users to write their own aliases for builds(<t>, populate_full_signature=True) such that they would have the desired type-inference behavior.

Some Implementation Details

This signature reflection occurs only when:

• populate_full_signature=True is specified
• and zen_partial=False
• and when no *args/**kwargs** are provided by the user to `builds
• and no "parents" have been supplied via builds_bases

The reason for this is that when any of these conditions do not hold, then the user does not need to specify all of the parameters in the target's signature. E.g. consider:

Conf = builds(f, x=1, populate_full_signature=True)

because x=1 is specified, the true signature of Conf is (x: int = 1. y: str = "hello"). Unfortunately, there is not way (that I know of) to tell the type-checkers about this modification to the signature. Thus if we enabled signature-reflection in this case, then you would get the false-negative:

# This is OK at runtime, but
# type-checker sees error: missing value x
Conf()

I do not want users to ever encounter false-positives due to our type annotations. Thus we disable signature reflection here:

It is also noteworthy that functions with overloads are not amenable to being ParamSpec'd:

[rsokl]

VSCode even reflects the docstring!

[esafak]

Is this only for Python 3.10+?

[rsokl]

Is this only for Python 3.10+?

No, this works for Python 3.6+. We use the backport of ParamSpec provided by typing_extensions.

[rsokl]

Note: I updated the original post substantially to reflect the fact that make_custom_builds_fn now has "dynamic" overloads. Thus this is no longer on the To-Do list!

[rsokl]

One unfortunate thing is that PyCharm can't keep up with almost any of these annotations.. It's type-checker is not good enough :/

[rsokl]

@jgbos this is a pretty big PR, but I tried to keep the commits pretty modular and tidy. When you have the chance to review this, I recommend doing so commit-by-commit, rather than scrolling through all of the diffs by themselves.