From 69f7603f09902500d454e06a4746ae31c9cc4211 Mon Sep 17 00:00:00 2001 From: Sveinung Gundersen Date: Fri, 31 Mar 2023 15:52:16 +0200 Subject: [PATCH] Temp commit trying out fixes to provide autocomplete on the job_func callable --- .../api/protocols/private/compute/job.py | 67 ++++++++++--------- src/omnipy/api/protocols/private/util.py | 9 ++- src/omnipy/api/protocols/public/compute.py | 30 +++++---- src/omnipy/compute/func_job.py | 12 ++-- src/omnipy/compute/job.py | 40 +++++++---- src/omnipy/compute/task.py | 64 ++++++++++++++---- src/omnipy/modules/fairtracks/tasks.py | 2 +- src/omnipy/util/callable_decorator_cls.py | 26 ++++--- 8 files changed, 164 insertions(+), 86 deletions(-) diff --git a/src/omnipy/api/protocols/private/compute/job.py b/src/omnipy/api/protocols/private/compute/job.py index 18073dfb..bd094121 100644 --- a/src/omnipy/api/protocols/private/compute/job.py +++ b/src/omnipy/api/protocols/private/compute/job.py @@ -2,7 +2,7 @@ from datetime import datetime from types import MappingProxyType -from typing import Any, Callable, Dict, Mapping, Optional, Protocol, Tuple, Type +from typing import Any, Callable, Dict, Mapping, Optional, Protocol, Tuple, Type, TypeVar from omnipy.api.enums import PersistOutputsOptions, RestoreOutputsOptions from omnipy.api.protocols.private.compute.job_creator import IsJobCreator @@ -17,8 +17,10 @@ TaskTemplateCovT, TaskTemplateT) +C = TypeVar('C', bound=Callable) -class IsJobBase(CanLog, IsUniquelyNamedJob, Protocol): + +class IsJobBase(CanLog, IsUniquelyNamedJob, Protocol[C]): """""" @property def _job_creator(self) -> IsJobCreator: @@ -40,57 +42,54 @@ def __eq__(self, other: object): ... @classmethod - def _create_job_template(cls, *args: object, **kwargs: object) -> IsJobTemplate: + def _create_job_template(cls, *args: object, **kwargs: object) -> IsJobTemplate[C]: ... @classmethod - def _create_job(cls, *args: object, **kwargs: object) -> IsJob: + def _create_job(cls, *args: object, **kwargs: object) -> IsJob[C]: ... - def _apply(self) -> IsJob: + def _apply(self) -> IsJob[C]: ... - def _refine(self, *args: Any, update: bool = True, **kwargs: object) -> IsJobTemplate: + def _refine(self, *args: Any, update: bool = True, **kwargs: object) -> IsJobTemplate[C]: ... - def _revise(self) -> IsJobTemplate: + def _revise(self) -> IsJobTemplate[C]: ... - def _call_job_template(self, *args: object, **kwargs: object) -> object: - ... + _call_job_template: C + _call_job: C - def _call_job(self, *args: object, **kwargs: object) -> object: - ... +class IsJobBaseCallable(IsJobBase[C], Protocol[C]): + __call__: C -class IsJob(IsJobBase, Protocol): + +class IsJob(IsJobBaseCallable[C], Protocol[C]): """""" @property def time_of_cur_toplevel_flow_run(self) -> Optional[datetime]: ... @classmethod - def create_job(cls, *args: object, **kwargs: object) -> IsJob: - ... - - def __call__(self, *args: object, **kwargs: object) -> object: + def create_job(cls, *args: object, **kwargs: object) -> IsJob[C]: ... def _apply_engine_decorator(self, engine: IsEngine) -> None: ... -class IsJobTemplate(IsJobBase, Protocol): +class IsJobTemplate(IsJobBaseCallable[C], Protocol[C]): """""" @classmethod - def create_job_template(cls, *args: object, **kwargs: object) -> IsJobTemplate: + def create_job_template(cls, *args: object, **kwargs: object) -> IsJobTemplate[C]: ... - def run(self, *args: object, **kwargs: object) -> object: - ... + run: C -class IsFuncArgJobBase(IsJob, Protocol): +class IsFuncArgJobBase(IsJob[C], Protocol[C]): """""" @property def param_signatures(self) -> MappingProxyType: @@ -139,6 +138,7 @@ def get_call_args(self, *args: object, **kwargs: object) -> Dict[str, object]: ... +## Change? class IsPlainFuncArgJobBase(Protocol): """""" _job_func: Callable @@ -147,13 +147,16 @@ def _accept_call_func_decorator(self, call_func_decorator: GeneralDecorator) -> ... -class IsFuncArgJob(IsFuncArgJobBase, Protocol[JobT]): +class IsFuncArgJob(IsFuncArgJobBase[C], Protocol[JobT, C]): """""" def revise(self) -> JobT: ... -class IsFuncArgJobTemplateCallable(Protocol[JobTemplateT]): +CC = TypeVar('CC', bound=Callable, contravariant=True) + + +class IsFuncArgJobTemplateCallable(Protocol[JobTemplateT, CC]): """""" def __call__( self, @@ -165,11 +168,11 @@ def __call__( fixed_params: Optional[Mapping[str, object]] = None, param_key_map: Optional[Mapping[str, str]] = None, **kwargs: object, - ) -> Callable[[Callable], JobTemplateT]: + ) -> Callable[[CC], JobTemplateT]: ... -class IsFuncArgJobTemplate(IsJobTemplate, IsFuncArgJobBase, Protocol[JobTemplateT, JobT]): +class IsFuncArgJobTemplate(IsJobTemplate[C], IsFuncArgJobBase[C], Protocol[JobTemplateT, JobT, C]): """""" def refine(self, *args: Any, @@ -188,16 +191,16 @@ def apply(self) -> JobT: ... -class IsTaskTemplateArgsJobBase(IsFuncArgJobBase, Protocol[TaskTemplateCovT]): +class IsTaskTemplateArgsJobBase(IsFuncArgJobBase[C], Protocol[TaskTemplateCovT, C]): """""" @property def task_templates(self) -> Tuple[TaskTemplateCovT, ...]: ... -class IsTaskTemplateArgsJob(IsTaskTemplateArgsJobBase[TaskTemplateCovT], - IsFuncArgJob[JobT], - Protocol[TaskTemplateCovT, JobT]): +class IsTaskTemplateArgsJob(IsTaskTemplateArgsJobBase[TaskTemplateCovT, C], + IsFuncArgJob[JobT, C], + Protocol[TaskTemplateCovT, JobT, C]): """""" @@ -218,9 +221,9 @@ def __call__( ... -class IsTaskTemplateArgsJobTemplate(IsFuncArgJobTemplate[JobTemplateT, JobT], - IsTaskTemplateArgsJobBase[TaskTemplateT], - Protocol[TaskTemplateT, JobTemplateT, JobT]): +class IsTaskTemplateArgsJobTemplate(IsFuncArgJobTemplate[JobTemplateT, JobT, C], + IsTaskTemplateArgsJobBase[TaskTemplateT, C], + Protocol[TaskTemplateT, JobTemplateT, JobT, C]): """""" def refine(self, *task_templates: TaskTemplateT, diff --git a/src/omnipy/api/protocols/private/util.py b/src/omnipy/api/protocols/private/util.py index 5aaeccdf..50ef726b 100644 --- a/src/omnipy/api/protocols/private/util.py +++ b/src/omnipy/api/protocols/private/util.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable, Protocol, runtime_checkable +from typing import Callable, Protocol, runtime_checkable, TypeVar from omnipy.api.types import DecoratorClassT @@ -12,8 +12,11 @@ def __call__(self, callable_arg: Callable, /, *args: object, **kwargs: object) - ... +CC = TypeVar('CC', bound=Callable, contravariant=True) + + @runtime_checkable -class IsCallableClass(Protocol[DecoratorClassT]): +class IsCallableClass(Protocol[DecoratorClassT, CC]): """""" - def __call__(self, *args: object, **kwargs: object) -> Callable[[Callable], DecoratorClassT]: + def __call__(self, *args: object, **kwargs: object) -> Callable[[CC], DecoratorClassT]: ... diff --git a/src/omnipy/api/protocols/public/compute.py b/src/omnipy/api/protocols/public/compute.py index 965c16b0..4b273065 100644 --- a/src/omnipy/api/protocols/public/compute.py +++ b/src/omnipy/api/protocols/public/compute.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import datetime -from typing import Optional, Protocol +from typing import Callable, Optional, Protocol, TypeVar from omnipy.api.protocols.private.compute.job import (IsFuncArgJob, IsFuncArgJobTemplate, @@ -9,13 +9,15 @@ IsTaskTemplateArgsJobTemplate) from omnipy.api.protocols.private.compute.mixins import IsNestedContext +C = TypeVar('C', bound=Callable) -class IsTaskTemplate(IsFuncArgJobTemplate['IsTaskTemplate', 'IsTask'], Protocol): + +class IsTaskTemplate(IsFuncArgJobTemplate['IsTaskTemplate', 'IsTask', C], Protocol[C]): """""" ... -class IsTask(IsFuncArgJob[IsTaskTemplate], Protocol): +class IsTask(IsFuncArgJob[IsTaskTemplate, C], Protocol[C]): """""" ... @@ -38,39 +40,43 @@ def time_of_last_run(self) -> Optional[datetime]: class IsLinearFlowTemplate(IsTaskTemplateArgsJobTemplate[IsTaskTemplate, 'IsLinearFlowTemplate', - 'IsLinearFlow'], + 'IsLinearFlow', + C], IsFlowTemplate, - Protocol): + Protocol[C]): """""" ... -class IsLinearFlow(IsTaskTemplateArgsJob[IsTaskTemplate, IsLinearFlowTemplate], IsFlow, Protocol): +class IsLinearFlow(IsTaskTemplateArgsJob[IsTaskTemplate, IsLinearFlowTemplate, C], + IsFlow, + Protocol[C]): """""" ... class IsDagFlowTemplate(IsTaskTemplateArgsJobTemplate[IsTaskTemplate, 'IsDagFlowTemplate', - 'IsDagFlow'], + 'IsDagFlow', + C], IsFlowTemplate, - Protocol): + Protocol[C]): """""" ... -class IsDagFlow(IsTaskTemplateArgsJob[IsTaskTemplate, IsDagFlowTemplate], IsFlow, Protocol): +class IsDagFlow(IsTaskTemplateArgsJob[IsTaskTemplate, IsDagFlowTemplate, C], IsFlow, Protocol[C]): """""" ... -class IsFuncFlowTemplate(IsFuncArgJobTemplate['IsFuncFlowTemplate', 'IsFuncFlow'], +class IsFuncFlowTemplate(IsFuncArgJobTemplate['IsFuncFlowTemplate', 'IsFuncFlow', C], IsFlowTemplate, - Protocol): + Protocol[C]): """""" ... -class IsFuncFlow(IsFuncArgJob[IsFuncFlowTemplate], Protocol): +class IsFuncFlow(IsFuncArgJob[IsFuncFlowTemplate, C], Protocol[C]): """""" ... diff --git a/src/omnipy/compute/func_job.py b/src/omnipy/compute/func_job.py index 45a67563..f9ae5c59 100644 --- a/src/omnipy/compute/func_job.py +++ b/src/omnipy/compute/func_job.py @@ -1,5 +1,5 @@ import asyncio -from typing import Callable, Tuple +from typing import Callable, Generic, Tuple, TypeVar from omnipy.api.types import GeneralDecorator from omnipy.compute.job import JobBase @@ -9,12 +9,14 @@ from omnipy.compute.mixins.result_key import ResultKeyFuncJobBaseMixin from omnipy.compute.mixins.serialize import SerializerFuncJobBaseMixin +C = TypeVar('C', bound=Callable) -class PlainFuncArgJobBase(JobBase): - def __init__(self, job_func: Callable, *args: object, **kwargs: object) -> None: + +class PlainFuncArgJobBase(JobBase, Generic[C]): + def __init__(self, job_func: C, *args: object, **kwargs: object) -> None: self._job_func = job_func - def _get_init_args(self) -> Tuple[object, ...]: + def _get_init_args(self) -> Tuple[C]: return self._job_func, def has_coroutine_func(self) -> bool: @@ -37,7 +39,7 @@ def _accept_call_func_decorator(self, call_func_decorator: GeneralDecorator) -> # Extra level needed for mixins to be able to overload _call_job (and possibly other methods) -class FuncArgJobBase(PlainFuncArgJobBase): +class FuncArgJobBase(PlainFuncArgJobBase[C], Generic[C]): ... diff --git a/src/omnipy/compute/job.py b/src/omnipy/compute/job.py index 6ddcbe15..c4c03981 100644 --- a/src/omnipy/compute/job.py +++ b/src/omnipy/compute/job.py @@ -5,7 +5,17 @@ from functools import update_wrapper import logging from types import MappingProxyType -from typing import Any, cast, Dict, Hashable, Optional, Tuple, Type, Union +from typing import (Any, + Callable, + cast, + Dict, + Generic, + Hashable, + Optional, + Tuple, + Type, + TypeVar, + Union) from omnipy.api.exceptions import JobStateException from omnipy.api.protocols.private.compute.job import IsJob, IsJobBase, IsJobTemplate @@ -176,14 +186,18 @@ def _check_engine(self, engine_protocol: Type): f'job runner protocol: {engine_protocol.__name__}') -class JobTemplateMixin: +C = TypeVar('C', bound=Callable) + + +class JobTemplateMixin(Generic[C]): def __init__(self, *args, **kwargs): if JobBase not in self.__class__.__mro__: raise TypeError('JobTemplateMixin is not meant to be instantiated outside the context ' 'of a JobBase subclass.') @classmethod - def create_job_template(cls: Type[IsJobBase], *args: object, **kwargs: object) -> IsJobTemplate: + def create_job_template(cls: Type[IsJobBase[C]], *args: object, + **kwargs: object) -> IsJobTemplate[C]: return cls._create_job_template(*args, **kwargs) def run(self, *args: object, **kwargs: object) -> object: @@ -191,22 +205,22 @@ def run(self, *args: object, **kwargs: object) -> object: return self.apply()(*args, **kwargs) - def apply(self) -> IsJob: - self_as_job_base = cast(IsJobBase, self) + def apply(self) -> IsJob[C]: + self_as_job_base = cast(IsJobBase[C], self) job = self_as_job_base._apply() update_wrapper(job, self, updated=[]) return job - def refine(self, *args: Any, update: bool = True, **kwargs: object) -> IsJobTemplate: - self_as_job_base = cast(IsJobBase, self) + def refine(self, *args: Any, update: bool = True, **kwargs: object) -> IsJobTemplate[C]: + self_as_job_base = cast(IsJobBase[C], self) return self_as_job_base._refine(*args, update=update, **kwargs) def __call__(self, *args: object, **kwargs: object) -> object: - self_as_job_base = cast(IsJobBase, self) + self_as_job_base = cast(IsJobBase[C], self) return self_as_job_base._call_job_template(*args, **kwargs) -class JobMixin(DynamicMixinAcceptor): +class JobMixin(DynamicMixinAcceptor, Generic[C]): def __init__(self, *args, **kwargs): if JobBase not in self.__class__.__mro__: raise TypeError('JobMixin is not meant to be instantiated outside the context ' @@ -218,22 +232,22 @@ def _apply_engine_decorator(self, engine: IsEngine) -> None: @property def time_of_cur_toplevel_flow_run(self) -> Optional[datetime]: - self_as_job_base = cast(IsJobBase, self) + self_as_job_base = cast(IsJobBase[C], self) return self_as_job_base._job_creator.time_of_cur_toplevel_nested_context_run @classmethod def create_job(cls, *args: object, **kwargs: object) -> IsJob: - cls_as_job_base = cast(IsJobBase, cls) + cls_as_job_base = cast(IsJobBase[C], cls) return cls_as_job_base._create_job(*args, **kwargs) def revise(self) -> IsJobTemplate: - self_as_job_base = cast(IsJobBase, self) + self_as_job_base = cast(IsJobBase[C], self) job_template = self_as_job_base._revise() update_wrapper(job_template, self, updated=[]) return job_template def __call__(self, *args: object, **kwargs: object) -> object: - self_as_job_base = cast(IsJobBase, self) + self_as_job_base = cast(IsJobBase[C], self) try: return self_as_job_base._call_job(*args, **kwargs) diff --git a/src/omnipy/compute/task.py b/src/omnipy/compute/task.py index 51ad4b8b..f6c52c50 100644 --- a/src/omnipy/compute/task.py +++ b/src/omnipy/compute/task.py @@ -1,21 +1,30 @@ from __future__ import annotations -from typing import cast, Type +from typing import Callable, cast, Generic, Protocol, Type, TypeVar from omnipy.api.protocols.private.compute.job import (IsFuncArgJobTemplateCallable, IsJob, IsJobTemplate) from omnipy.api.protocols.private.engine import IsEngine +from omnipy.api.protocols.private.util import IsCallableClass from omnipy.api.protocols.public.compute import IsTask, IsTaskTemplate from omnipy.api.protocols.public.engine import IsTaskRunnerEngine +from omnipy.api.types import DecoratorClassT from omnipy.compute.func_job import FuncArgJobBase from omnipy.compute.job import JobMixin, JobTemplateMixin -from omnipy.util.callable_decorator_cls import callable_decorator_cls +from omnipy.util.callable_decorator_cls import callable_decorator_cls, IsDecoratorClass + +C = TypeVar('C', bound=Callable) def task_template_callable_decorator_cls( - cls: Type[TaskTemplate]) -> IsFuncArgJobTemplateCallable[IsTaskTemplate]: - return cast(IsFuncArgJobTemplateCallable[IsTaskTemplate], callable_decorator_cls(cls)) + cls: Type[TaskTemplateInner[C]]) -> IsFuncArgJobTemplateCallable[IsTaskTemplate[C], C]: + decorated = callable_decorator_cls(cls) + reveal_type(decorated) + return cast(IsFuncArgJobTemplateCallable[IsTaskTemplate[C], C], decorated) + + +reveal_type(task_template_callable_decorator_cls) class TaskBase: @@ -23,24 +32,55 @@ class TaskBase: ... -@task_template_callable_decorator_cls -class TaskTemplate(JobTemplateMixin, TaskBase, FuncArgJobBase): +# @task_template_callable_decorator_cls +class TaskTemplateInner(FuncArgJobBase[C], + JobTemplateMixin[C], + TaskBase, + Generic[C], + IsDecoratorClass[C]): """""" @classmethod - def _get_job_subcls_for_apply(cls) -> Type[IsJob]: - return cast(Type[IsTask], Task) + def _get_job_subcls_for_apply(cls) -> Type[IsJob[C]]: + return cast(Type[IsTask[C]], Task[C]) + + +# +# class A(Protocol): +# def __call__(self, a: int) -> int: +# ... + + +def a(a: int) -> int: + return a + 2 + + +# +# def get_task_template() -> IsFuncArgJobTemplateCallable[IsTaskTemplate[C], C]: +# return task_template_callable_decorator_cls(TaskTemplateInner[C]) +# +# +# TaskTemplate = get_task_template() + +TaskTemplate = task_template_callable_decorator_cls(TaskTemplateInner) +reveal_type(TaskTemplate) + +tstt = TaskTemplate() +reveal_type(tstt) + +t = tstt(a) +reveal_type(t) -class Task(JobMixin, TaskBase, FuncArgJobBase): +class Task(JobMixin[C], TaskBase, FuncArgJobBase[C], Generic[C]): def _apply_engine_decorator(self, engine: IsEngine) -> None: if self.engine: engine = cast(IsTaskRunnerEngine, self.engine) - self_with_mixins = cast(IsTask, self) + self_with_mixins = cast(IsTask[C], self) engine.apply_task_decorator(self_with_mixins, self._accept_call_func_decorator) @classmethod - def _get_job_template_subcls_for_revise(cls) -> Type[IsJobTemplate]: - return cast(Type[IsTaskTemplate], TaskTemplate) + def _get_job_template_subcls_for_revise(cls) -> Type[IsJobTemplate[C]]: + return cast(Type[IsTaskTemplate[C]], TaskTemplate[C]) # TODO: Would we need the possibility to refine task templates by adding new task parameters? diff --git a/src/omnipy/modules/fairtracks/tasks.py b/src/omnipy/modules/fairtracks/tasks.py index 469032e8..04ce74d8 100644 --- a/src/omnipy/modules/fairtracks/tasks.py +++ b/src/omnipy/modules/fairtracks/tasks.py @@ -9,7 +9,7 @@ from .functions import encode_api -@mypy_fix_task_template +# @mypy_fix_task_template @TaskTemplate() def import_dataset_from_encode(endpoints: Iterable[constr(min_length=1)], max_data_item_count: PositiveInt) -> JsonDataset: diff --git a/src/omnipy/util/callable_decorator_cls.py b/src/omnipy/util/callable_decorator_cls.py index 26e982ae..ff8ae031 100644 --- a/src/omnipy/util/callable_decorator_cls.py +++ b/src/omnipy/util/callable_decorator_cls.py @@ -1,12 +1,22 @@ from functools import update_wrapper from types import MethodWrapperType -from typing import Callable, cast, Type +from typing import Callable, cast, Protocol, Type, TypeVar from omnipy.api.protocols.private.util import IsCallableClass, IsCallableParamAfterSelf from omnipy.api.types import DecoratorClassT +CC = TypeVar('CC', bound=Callable, contravariant=True) -def callable_decorator_cls(cls: Type[DecoratorClassT]) -> IsCallableClass[DecoratorClassT]: +C = TypeVar('C', bound=Callable) + + +class IsDecoratorClass(Protocol[CC]): + def __call__(self, func: CC, *args: object, **kwargs: object) -> None: + ... + + +def callable_decorator_cls( + cls: Type[IsDecoratorClass[CC]]) -> IsCallableClass[IsDecoratorClass[CC], CC]: """ "Meta-decorator" that allows any class to function as a decorator for a callable. @@ -20,7 +30,7 @@ def callable_decorator_cls(cls: Type[DecoratorClassT]) -> IsCallableClass[Decora cls._wrapped_call: Callable = cast(Callable, cls.__call__) def _forward_call_to_obj_if_callable(self, *args: object, - **kwargs: object) -> Type[DecoratorClassT]: + **kwargs: object) -> Type[IsDecoratorClass[CC]]: """ __call__ method at the class level which forward the call to instance-level call methods, if present (hardcoded as '_obj_call()'). This is needed due to the peculiarity that Python @@ -56,7 +66,7 @@ def _real_callable(arg: object) -> bool: _wrapped_new: Callable = cls.__new__ - def _new_wrapper(cls, *args: object, **kwargs: object) -> DecoratorClassT: + def _new_wrapper(cls, *args: object, **kwargs: object) -> IsDecoratorClass[CC]: if _wrapped_new is object.__new__: obj = _wrapped_new(cls) else: @@ -70,13 +80,13 @@ def _new_wrapper(cls, *args: object, **kwargs: object) -> DecoratorClassT: def _init_wrapper(self, *args: object, **kwargs: object) -> None: args_list = list(args) - def _init(callable_arg: Callable) -> None: + def _init(callable_arg: CC) -> None: _wrapped_init(self, callable_arg, *args_list, **kwargs) update_wrapper(self, callable_arg, updated=[]) if len(args_list) == 1 and _real_callable(args_list[0]): # Decorate the callable directly - _callable_arg: Callable = cast(Callable, args_list[0]) + _callable_arg = cast(CC, args_list[0]) args_list.pop(0) _init(_callable_arg) else: @@ -84,7 +94,7 @@ def _init(callable_arg: Callable) -> None: # class-level __call__ method. When this method is called, the provided _callable_arg # is decorated. def _init_as_obj_call_method( - self, _callable_arg: Callable) -> Type[DecoratorClassT]: # noqa + self, _callable_arrg: CC) -> Type[IsDecoratorClass[CC]]: # noqa _init(_callable_arg) del self._obj_call return self @@ -99,4 +109,4 @@ def _init_as_obj_call_method( setattr(cls, '__new__', _new_wrapper) - return cast(IsCallableClass[DecoratorClassT], cls) + return cast(IsCallableClass[IsDecoratorClass[CC], CC], cls)