Coverage for src / invariant / traits.py: 93.75%

1"""Operation execution traits. 

2 

3Traits are portable op metadata. Executors and schedulers may use them to route 

4work, but traits are not part of artifact identity. 

5""" 

6 

7from collections.abc import Callable, Iterable 

8from enum import Enum 

9from typing import Any 

10 

11 

12class OpTrait(str, Enum): 

13 """Built-in operation execution traits.""" 

14 

15 BLOCKING = "blocking" 

16 IO_BOUND = "io-bound" 

17 CPU_BOUND = "cpu-bound" 

18 THREAD_SAFE = "thread-safe" 

19 PROCESS_SAFE = "process-safe" 

20 

21 

22TraitLike = str | Enum 

23 

24_OP_TRAITS_ATTR = "__invariant_op_traits__" 

25 

26 

27def normalize_trait(trait: TraitLike) -> str: 

28 """Normalize a trait enum or string to its wire/storage string.""" 

29 value = trait.value if isinstance(trait, Enum) else trait 

30 if not isinstance(value, str): 

31 value = str(value) 

32 if not value: 

33 raise ValueError("Operation trait cannot be empty") 

34 return value 

35 

36 

37def normalize_traits(traits: Iterable[TraitLike] | None = None) -> frozenset[str]: 

38 """Normalize a trait iterable to a stable set of strings.""" 

39 if traits is None: 

40 return frozenset() 

41 return frozenset(normalize_trait(trait) for trait in traits) 

42 

43 

44def op_traits(*traits: TraitLike) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

45 """Attach execution traits to an operation callable.""" 

46 normalized = normalize_traits(traits) 

47 

48 def decorate(fn: Callable[..., Any]) -> Callable[..., Any]: 

49 existing = normalize_traits(getattr(fn, _OP_TRAITS_ATTR, frozenset())) 

50 setattr(fn, _OP_TRAITS_ATTR, existing | normalized) 

51 return fn 

52 

53 return decorate 

54 

55 

56def decorated_traits(op: Callable[..., Any]) -> frozenset[str]: 

57 """Return traits attached by ``@op_traits``.""" 

58 return normalize_traits(getattr(op, _OP_TRAITS_ATTR, frozenset())) 

59 

60 

61__all__ = [ 

62 "OpTrait", 

63 "TraitLike", 

64 "decorated_traits", 

65 "normalize_trait", 

66 "normalize_traits", 

67 "op_traits", 

68]