diff --git a/src/breakshaft/convertor.py b/src/breakshaft/convertor.py index 8688e2a..9439798 100644 --- a/src/breakshaft/convertor.py +++ b/src/breakshaft/convertor.py @@ -1,10 +1,12 @@ from __future__ import annotations -from typing import Optional, Callable, Unpack, TypeVarTuple, TypeVar, Awaitable, Any, Sequence + +import collections.abc +from typing import Optional, Callable, Unpack, TypeVarTuple, TypeVar, Awaitable, Any, Sequence, Iterable from .graph_walker import GraphWalker from .models import ConversionPoint, Callgraph from .renderer import ConvertorRenderer, InTimeGenerationConvertorRenderer -from .util import extract_return_type +from .util import extract_return_type, universal_qualname Tin = TypeVarTuple('Tin') Tout = TypeVar('Tout') @@ -69,8 +71,11 @@ class ConvRepo: def convertor_set(self): return self._convertor_set - def add_injector(self, func: Callable, rettype: Optional[type] = None): - self._convertor_set |= set(ConversionPoint.from_fn(func, rettype=rettype)) + def add_injector(self, + func: Callable, + rettype: Optional[type] = None, + type_remap: Optional[dict[str, type]] = None): + self._convertor_set |= set(ConversionPoint.from_fn(func, rettype=rettype, type_remap=type_remap)) def _callseq_from_callgraph(self, cg: Callgraph) -> list[ConversionPoint]: if len(cg.variants) == 0: @@ -97,12 +102,12 @@ class ConvRepo: def get_callseq(self, injectors: frozenset[ConversionPoint], from_types: frozenset[type], - fn: Callable, + fn: Callable | Iterable[ConversionPoint] | ConversionPoint, force_commutative: bool) -> list[ConversionPoint]: cg = self.walker.generate_callgraph(injectors, from_types, fn) if cg is None: - raise ValueError(f'Unable to compute conversion graph on {from_types}->{fn.__qualname__}') + raise ValueError(f'Unable to compute conversion graph on {from_types}->{universal_qualname(fn)}') exploded = self.walker.explode_callgraph_branches(cg, from_types) @@ -116,14 +121,22 @@ class ConvRepo: callseq = self._callseq_from_callgraph(Callgraph(frozenset([selected[0]]))) if len(callseq) > 0: - injects = extract_return_type(fn) + injects = None + if isinstance(fn, collections.abc.Iterable): + for f in fn: + injects = f.injects + break + elif isinstance(fn, ConversionPoint): + injects = fn.injects + else: + injects = extract_return_type(fn) callseq[-1] = callseq[-1].copy_with(injects=injects) return callseq def get_conversion(self, from_types: Sequence[type[Unpack[Tin]]], - fn: Callable[..., Tout], + fn: Callable[..., Tout] | Iterable[ConversionPoint] | ConversionPoint, force_commutative: bool = True, allow_async: bool = True, allow_sync: bool = True, @@ -138,9 +151,9 @@ class ConvRepo: setattr(ret_fn, '__breakshaft_callseq__', callseq) return ret_fn - def mark_injector(self, *, rettype: Optional[type] = None): + def mark_injector(self, *, rettype: Optional[type] = None, type_remap: Optional[dict[str, type]] = None): def inner(func: Callable): - self.add_injector(func) + self.add_injector(func, rettype=rettype, type_remap=type_remap) return func return inner diff --git a/src/breakshaft/graph_walker.py b/src/breakshaft/graph_walker.py index b7ff4aa..f1e2d23 100644 --- a/src/breakshaft/graph_walker.py +++ b/src/breakshaft/graph_walker.py @@ -1,9 +1,11 @@ +import collections.abc import typing from types import NoneType from typing import Callable, Optional from .models import ConversionPoint, Callgraph, CallgraphVariant, TransformationPoint, CompositionDirection from .util import extract_func_argtypes, all_combinations, extract_func_argtypes_seq, extract_return_type +from typing import Iterable class GraphWalker: @@ -12,17 +14,25 @@ class GraphWalker: def generate_callgraph(cls, injectors: frozenset[ConversionPoint], from_types: frozenset[type], - consumer_fn: Callable) -> Optional[Callgraph]: + consumer_fn: Callable | Iterable[ConversionPoint] | ConversionPoint) -> Optional[Callgraph]: branches: frozenset[Callgraph] = frozenset() - rettype = extract_return_type(consumer_fn) # Хак, чтобы вынудить систему поставить первым преобразованием требуемый consumer # Новый TypeAliasType каждый раз будет иметь эксклюзивный хэш, вне зависимости от содержимого # При этом, TypeAliasType также выступает в роли ключа преобразования # Это позволяет переложить обработку аргументов consumer на внутренние механизмы построения графа преобразований type _tmp_type_for_consumer = object - injectors |= set(ConversionPoint.from_fn(consumer_fn, _tmp_type_for_consumer)) + + if isinstance(consumer_fn, collections.abc.Iterable): + new_consumer_injectors = set() + for fn in consumer_fn: + new_consumer_injectors.add(fn.copy_with(injects=_tmp_type_for_consumer)) + injectors |= new_consumer_injectors + elif isinstance(consumer_fn, ConversionPoint): + injectors |= set(consumer_fn.copy_with(injects=_tmp_type_for_consumer)) + else: + injectors |= set(ConversionPoint.from_fn(consumer_fn, _tmp_type_for_consumer)) return cls.generate_callgraph_singletype(injectors, from_types, _tmp_type_for_consumer) @@ -143,7 +153,7 @@ class GraphWalker: if len(variants) > 1: # sorting by first injector func name for creating minimal cosistancy # could lead to heizenbugs due to incosistancy in path selection between calls - variants.sort(key=lambda x: x.injector.fn.__qualname__) + variants.sort(key=lambda x: universal_qualname(x.injector.fn)) return variants if len(variants) < 2: diff --git a/src/breakshaft/models.py b/src/breakshaft/models.py index abd4b2a..8aab9e0 100644 --- a/src/breakshaft/models.py +++ b/src/breakshaft/models.py @@ -11,7 +11,7 @@ from typing import Callable, Optional, get_type_hints, get_origin, Generator, ge from .util import extract_func_argtypes, extract_func_argtypes_seq, is_sync_context_manager_factory, \ is_async_context_manager_factory, \ all_combinations, is_context_manager_factory, extract_func_arg_defaults, extract_func_args, extract_func_argnames, \ - get_tuple_types, is_basic_type_annot + get_tuple_types, is_basic_type_annot, universal_qualname @dataclass(frozen=True) @@ -34,15 +34,8 @@ class ConversionPoint: return hash((self.fn, self.injects, self.requires)) def __repr__(self): - if '__qualname__' in dir(self.injects): - injects_name = self.injects.__qualname__ - else: - injects_name = str(self.injects) - - if '__qualname__' in dir(self.fn): - fn_name = self.fn.__qualname__ - else: - fn_name = str(self.fn) + injects_name = universal_qualname(self.injects) + fn_name = universal_qualname(self.fn) return f'({",".join(map(str, self.requires))}) -> {injects_name}: {fn_name}' @@ -60,9 +53,15 @@ class ConversionPoint: return inspect.iscoroutinefunction(self.fn) or is_async_context_manager_factory(self.fn) @classmethod - def from_fn(cls, func: Callable, rettype: Optional[type] = None) -> list[ConversionPoint]: + def from_fn(cls, + func: Callable, + rettype: Optional[type] = None, + type_remap: Optional[dict[str, type]] = None) -> list[ConversionPoint]: + if type_remap is None: + annot = get_type_hints(func) + else: + annot = type_remap - annot = get_type_hints(func) fn_rettype = annot.get('return') if rettype is None: rettype = fn_rettype @@ -97,10 +96,10 @@ class ConversionPoint: if len(tuple_unwrapped) > 0 and Ellipsis not in tuple_unwrapped: for t in tuple_unwrapped: if not is_basic_type_annot(t): - ret += ConversionPoint.from_fn(func, rettype=t) + ret += ConversionPoint.from_fn(func, rettype=t, type_remap=type_remap) argtypes: list[list[type]] = [] - orig_args = extract_func_args(func) + orig_args = extract_func_args(func, type_remap) defaults = extract_func_arg_defaults(func) orig_argtypes = [] diff --git a/src/breakshaft/util.py b/src/breakshaft/util.py index deba74b..b0a0938 100644 --- a/src/breakshaft/util.py +++ b/src/breakshaft/util.py @@ -19,9 +19,13 @@ def extract_return_type(func: Callable) -> Optional[type]: return hints.get('return') -def extract_func_args(func: Callable) -> list[tuple[str, type]]: +def extract_func_args(func: Callable, type_hints_remap: Optional[dict[str, type]] = None) -> list[tuple[str, type]]: sig = inspect.signature(func) - type_hints = get_type_hints(func) + if type_hints_remap is None: + type_hints = get_type_hints(func) + else: + type_hints = type_hints_remap + params = sig.parameters args_info = [] @@ -127,3 +131,12 @@ def is_basic_type_annot(type_annot) -> bool: return all(is_basic_type_annot(arg) for arg in args) return False + + +def universal_qualname(any: Any) -> str: + if hasattr(any, '__qualname__'): + return any.__qualname__ + if hasattr(any, '__name__'): + return any.__name__ + + return str(any) diff --git a/tests/test_typehints_remap.py b/tests/test_typehints_remap.py new file mode 100644 index 0000000..0c6d4d0 --- /dev/null +++ b/tests/test_typehints_remap.py @@ -0,0 +1,47 @@ +from dataclasses import dataclass +from typing import Annotated + +import pytest + +from breakshaft.models import ConversionPoint +from src.breakshaft.convertor import ConvRepo + + +@dataclass +class A: + a: int + + +@dataclass +class B: + b: float + + +def test_basic(): + repo = ConvRepo() + + @repo.mark_injector() + def int_to_a(i: int) -> A: + return A(i) + + type HackInt = int + + def consumer(dep: A) -> int: + return dep.a + + type NewA = A + type_remap = {'dep': NewA, 'return': Annotated[HackInt, 'fuck']} + + assert len(ConversionPoint.from_fn(consumer, type_remap=type_remap)) == 1 + + with pytest.raises(ValueError): + fn1 = repo.get_conversion((int,), ConversionPoint.from_fn(consumer, type_remap=type_remap), + force_commutative=True, force_async=False, allow_async=False) + + repo.mark_injector(type_remap={'i': int, 'return': NewA})(int_to_a) + + fn1 = repo.get_conversion((int,), ConversionPoint.from_fn(consumer, type_remap=type_remap), + force_commutative=True, force_async=False, allow_async=False) + + assert fn1(42) == 42 +