Implement entrypoint typing #3387
Conversation
|
|
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
simonpernas
force-pushed
the
implement-entrypoint-typing
branch
from
5ca7e7c to
017381d
Compare
src/zenml/steps/base_step.py
Outdated
| @classmethod | ||
| def load_from_source(cls, source: Union[Source, str]) -> "BaseStep": | ||
| def load_from_source(cls, source: Union[Source, str]) -> "BaseStep[F]": |
There was a problem hiding this comment.
I think this should rather return BaseStep[Callable[..., Any]] which at least for me solved some type issues. As a more general question: Shouldn't we always use BaseStep[Callable[..., Any]] instead of BaseStep[Any] or is this already covered because F is bound to Callable[..., Any] anyway?
There was a problem hiding this comment.
With the version you commented on, if we tried to type a BaseStep with something that is not a Callable, like BaseStep[int], MyPy would report an error like this:
Value of type variable "F" of "BaseStep" cannot be "int" [type-var].
But it may still better to write BaseStep[Callable[..., Any]], at least for clarity. And anyway in the new version I get rid of this F entirely, and instead I use a ParamSpec as you suggested, giving typing like BaseStep[P, R], or if we don't care about the precise typing, BaseStep[..., Any]
src/zenml/steps/base_step.py
Outdated
|
|
||
|
|
||
| class BaseStep: | ||
| class _AbstractEntrypoint(Protocol[F]): |
There was a problem hiding this comment.
I'm guessing this was necessary to annotate the entrypoint? Do you think we could improve that even further by using https://docs.python.org/3/library/typing.html#typing.ParamSpec?
Without any testing, I was thinking of something like this:
P = ParamSpec("P")
R = TypeVar("R")
F = TypeVar("F", bound=Callable[P, R])
class BaseStep(Generic[F]):
@abstractmethod
def entrypoint(self, *args: P.args, **kwargs: P.kwargs) -> R:
passThere was a problem hiding this comment.
I did not think to use ParamSpec, but indeed it works here. Although not in the way you propose, because a TypeVar cannot be bound with another TypeVar (as far as I understand). But I think I came up with another cleaner solution that also addresses both your comments. Let me just verify that it actually works and passes MyPy tests
simonpernas
force-pushed
the
implement-entrypoint-typing
branch
from
017381d to
84aefb8
Compare
|
@simonpernas The CI is still failing, other than that looks good to me. Thanks for the PR! |
Add `from __future__ import annotations` to prevent `TypeError: 'type' object is not subscriptable` (for `BaseStep`) Change hints from `BaseStep[..., Any]` to `BaseStep[Any, Any]` when used for accessing the static method `load_from_source` (Ellipsis `...` used to not be acceptable for Generic types)
Describe changes
I implemented type hints for the entrypoint of functions decorated with
@step. With my changes,some_step.entrypointwill have the same typing information as the undecoratedsome_stepfunction. This is meant to help from a user point of view.I discussed briefly with Michael Schuster on Slack, because I do not have a proper testing environment for ZenML. He advised me to create a PR nonetheless, but I fear there may be issues that I missed because I could not run all the tests.
Pre-requisites
Please ensure you have done the following:
developand the open PR is targetingdevelop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.Types of changes