Client Lazy Loader

zenml.client_lazy_loader

Lazy loading functionality for Client methods.

ClientLazyLoader (BaseModel) pydantic-model

Lazy loader for Client methods.

Source code in zenml/client_lazy_loader.py

class ClientLazyLoader(BaseModel):
    """Lazy loader for Client methods."""

    method_name: str
    call_chain: List[_CallStep] = []

    _exclude_next_call: bool = PrivateAttr(False)

    def __getattr__(self, name: str) -> "ClientLazyLoader":
        """Get attribute not defined in ClientLazyLoader.

        Args:
            name: Name of the attribute to get.

        Returns:
            self
        """
        self_ = ClientLazyLoader(
            method_name=self.method_name, call_chain=self.call_chain.copy()
        )
        # workaround to protect from infinitely looping over in deepcopy called in invocations
        if name != "__deepcopy__":
            self_.call_chain.append(_CallStep(attribute_name=name))
        else:
            self_._exclude_next_call = True
        return self_

    def __call__(self, *args: Any, **kwargs: Any) -> "ClientLazyLoader":
        """Call mocked attribute.

        Args:
            args: Positional arguments.
            kwargs: Keyword arguments.

        Returns:
            self
        """
        # workaround to protect from infinitely looping over in deepcopy called in invocations
        if not self._exclude_next_call:
            self.call_chain.append(
                _CallStep(is_call=True, call_args=args, call_kwargs=kwargs)
            )
        self._exclude_next_call = False
        return self

    def __getitem__(self, item: Any) -> "ClientLazyLoader":
        """Get item from mocked attribute.

        Args:
            item: Item to get.

        Returns:
            self
        """
        self.call_chain.append(_CallStep(selector=item))
        return self

    def evaluate(self) -> Any:
        """Evaluate lazy loaded Client method.

        Returns:
            Evaluated lazy loader chain of calls.
        """
        from zenml.client import Client

        def _iterate_over_lazy_chain(
            self: "ClientLazyLoader", self_: Any, call_chain_: List[_CallStep]
        ) -> Any:
            next_step = call_chain_.pop(0)
            if next_step.is_call:
                self_ = self_(*next_step.call_args, **next_step.call_kwargs)
            elif next_step.selector:
                self_ = self_[next_step.selector]
            elif next_step.attribute_name:
                self_ = getattr(self_, next_step.attribute_name)
            else:
                raise RuntimeError(
                    f"Failed to evaluate lazy load chain `{self.method_name}` "
                    f"+ `{self.call_chain}`. Reach out to the ZenML team via "
                    "Slack or GitHub to check further."
                )
            return self_

        self_ = getattr(Client(), self.method_name)
        call_chain_ = self.call_chain.copy()
        while call_chain_:
            self_ = _iterate_over_lazy_chain(self, self_, call_chain_)
        return self_

__call__(self, *args, **kwargs) special

Call mocked attribute.

Parameters:

Name	Type	Description	Default
args	Any	Positional arguments.	()
kwargs	Any	Keyword arguments.	{}

Returns:

Type	Description
ClientLazyLoader	self

Source code in zenml/client_lazy_loader.py

def __call__(self, *args: Any, **kwargs: Any) -> "ClientLazyLoader":
    """Call mocked attribute.

    Args:
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        self
    """
    # workaround to protect from infinitely looping over in deepcopy called in invocations
    if not self._exclude_next_call:
        self.call_chain.append(
            _CallStep(is_call=True, call_args=args, call_kwargs=kwargs)
        )
    self._exclude_next_call = False
    return self

__getattr__(self, name) special

Get attribute not defined in ClientLazyLoader.

Parameters:

Name	Type	Description	Default
name	str	Name of the attribute to get.	required

Returns:

Type	Description
ClientLazyLoader	self

Source code in zenml/client_lazy_loader.py

def __getattr__(self, name: str) -> "ClientLazyLoader":
    """Get attribute not defined in ClientLazyLoader.

    Args:
        name: Name of the attribute to get.

    Returns:
        self
    """
    self_ = ClientLazyLoader(
        method_name=self.method_name, call_chain=self.call_chain.copy()
    )
    # workaround to protect from infinitely looping over in deepcopy called in invocations
    if name != "__deepcopy__":
        self_.call_chain.append(_CallStep(attribute_name=name))
    else:
        self_._exclude_next_call = True
    return self_

__getitem__(self, item) special

Get item from mocked attribute.

Parameters:

Name	Type	Description	Default
item	Any	Item to get.	required

Returns:

Type	Description
ClientLazyLoader	self

Source code in zenml/client_lazy_loader.py

def __getitem__(self, item: Any) -> "ClientLazyLoader":
    """Get item from mocked attribute.

    Args:
        item: Item to get.

    Returns:
        self
    """
    self.call_chain.append(_CallStep(selector=item))
    return self

evaluate(self)

Evaluate lazy loaded Client method.

Returns:

Type	Description
Any	Evaluated lazy loader chain of calls.

Source code in zenml/client_lazy_loader.py

def evaluate(self) -> Any:
    """Evaluate lazy loaded Client method.

    Returns:
        Evaluated lazy loader chain of calls.
    """
    from zenml.client import Client

    def _iterate_over_lazy_chain(
        self: "ClientLazyLoader", self_: Any, call_chain_: List[_CallStep]
    ) -> Any:
        next_step = call_chain_.pop(0)
        if next_step.is_call:
            self_ = self_(*next_step.call_args, **next_step.call_kwargs)
        elif next_step.selector:
            self_ = self_[next_step.selector]
        elif next_step.attribute_name:
            self_ = getattr(self_, next_step.attribute_name)
        else:
            raise RuntimeError(
                f"Failed to evaluate lazy load chain `{self.method_name}` "
                f"+ `{self.call_chain}`. Reach out to the ZenML team via "
                "Slack or GitHub to check further."
            )
        return self_

    self_ = getattr(Client(), self.method_name)
    call_chain_ = self.call_chain.copy()
    while call_chain_:
        self_ = _iterate_over_lazy_chain(self, self_, call_chain_)
    return self_

client_lazy_loader(method_name, *args, **kwargs)

Lazy loader for Client methods helper.

Usage:

def get_something(self, arg1: Any)->SomeResponse:
    if cll:=client_lazy_loader("get_something", arg1):
        return cll # type: ignore[return-value]
    return SomeResponse()

Parameters:

Name	Type	Description	Default
method_name	str	The name of the method to be called.	required
*args	Any	The arguments to be passed to the method.	()
**kwargs	Any	The keyword arguments to be passed to the method.	{}

Returns:

Type	Description
Optional[zenml.client_lazy_loader.ClientLazyLoader]	The result of the method call.

Source code in zenml/client_lazy_loader.py

def client_lazy_loader(
    method_name: str, *args: Any, **kwargs: Any
) -> Optional[ClientLazyLoader]:
    """Lazy loader for Client methods helper.

    Usage:
    ```
    def get_something(self, arg1: Any)->SomeResponse:
        if cll:=client_lazy_loader("get_something", arg1):
            return cll # type: ignore[return-value]
        return SomeResponse()
    ```

    Args:
        method_name: The name of the method to be called.
        *args: The arguments to be passed to the method.
        **kwargs: The keyword arguments to be passed to the method.

    Returns:
        The result of the method call.
    """
    from zenml import get_pipeline_context

    try:
        get_pipeline_context()
        cll = ClientLazyLoader(
            method_name=method_name,
        )
        return cll(*args, **kwargs)
    except RuntimeError:
        return None

evaluate_all_lazy_load_args_in_client_methods(cls)

Class wrapper to evaluate lazy loader arguments of all methods.

Parameters:

Name	Type	Description	Default
cls	Type[Client]	The class to wrap.	required

Returns:

Type	Description
Type[Client]	Wrapped class.

Source code in zenml/client_lazy_loader.py

def evaluate_all_lazy_load_args_in_client_methods(
    cls: Type["Client"],
) -> Type["Client"]:
    """Class wrapper to evaluate lazy loader arguments of all methods.

    Args:
        cls: The class to wrap.

    Returns:
        Wrapped class.
    """
    import inspect

    def _evaluate_args(func: Callable[..., Any]) -> Any:
        def _inner(*args: Any, **kwargs: Any) -> Any:
            is_instance_method = "self" in inspect.getfullargspec(func).args

            args_ = list(args)
            if not is_instance_method:
                from zenml.client import Client

                if args and isinstance(args[0], Client):
                    args_ = list(args[1:])

            for i in range(len(args_)):
                if isinstance(args_[i], dict):
                    with contextlib.suppress(ValueError):
                        args_[i] = ClientLazyLoader(**args_[i]).evaluate()
                elif isinstance(args_[i], ClientLazyLoader):
                    args_[i] = args_[i].evaluate()

            for k, v in kwargs.items():
                if isinstance(v, dict):
                    with contextlib.suppress(ValueError):
                        kwargs[k] = ClientLazyLoader(**v).evaluate()

            return func(*args_, **kwargs)

        return _inner

    def _decorate() -> Type["Client"]:
        for name, fn in inspect.getmembers(cls, inspect.isfunction):
            setattr(cls, name, _evaluate_args(fn))
        return cls

    return _decorate()