Skip to content

Commit 5ca7e7c

Browse files
committed
Remove type:ignore for entrypoint using a Protocol instead of abstractmethod
1 parent 25833d0 commit 5ca7e7c

File tree

1 file changed

+15
-16
lines changed

1 file changed

+15
-16
lines changed

src/zenml/steps/base_step.py

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import copy
1717
import hashlib
1818
import inspect
19-
from abc import abstractmethod
2019
from collections import defaultdict
2120
from typing import (
2221
TYPE_CHECKING,
@@ -27,6 +26,7 @@
2726
List,
2827
Mapping,
2928
Optional,
29+
Protocol,
3030
Sequence,
3131
Tuple,
3232
Type,
@@ -97,7 +97,20 @@
9797
F = TypeVar("F", bound=Callable[..., Any])
9898

9999

100-
class BaseStep(Generic[F]):
100+
class _AbstractEntrypoint(Protocol[F]):
101+
entrypoint: F
102+
"""Abstract method for core step logic.
103+
104+
Args:
105+
*args: Positional arguments passed to the step.
106+
**kwargs: Keyword arguments passed to the step.
107+
108+
Returns:
109+
The output of the step.
110+
"""
111+
112+
113+
class BaseStep(Generic[F], _AbstractEntrypoint[F]):
101114
"""Abstract base class for all ZenML steps."""
102115

103116
def __init__(
@@ -214,20 +227,6 @@ def __init__(
214227

215228
notebook_utils.try_to_save_notebook_cell_code(self.source_object)
216229

217-
@abstractmethod
218-
def entrypoint(self, *args: Any, **kwargs: Any) -> Any:
219-
"""Abstract method for core step logic.
220-
221-
Args:
222-
*args: Positional arguments passed to the step.
223-
**kwargs: Keyword arguments passed to the step.
224-
225-
Returns:
226-
The output of the step.
227-
"""
228-
229-
entrypoint: F # type:ignore[no-redef]
230-
231230
@classmethod
232231
def load_from_source(cls, source: Union[Source, str]) -> "BaseStep[F]":
233232
"""Loads a step from source.

0 commit comments

Comments
 (0)