Coverage for src/objectstate/parametric_axes.py: 93%

1""" 

2Parametric Axes Prototype - Arbitrary semantic axes for type construction. 

3 

4This module prototypes extending type() to support arbitrary axes beyond (B, S). 

5Proof-of-concept for PEP proposal: extending type() with axes= parameter. 

6 

7In the real CPython implementation, this logic would be in type.__new__ itself, 

8and every class would have __axes__ by default (empty MappingProxyType). 

9 

10This prototype uses AxesMeta as a stand-in for what type() would do natively. 

11Once any class uses metaclass=AxesMeta, all subclasses inherit it automatically. 

12 

13Usage: 

14 from objectstate.parametric_axes import AxesMeta 

15 

16 # One class in hierarchy uses metaclass=AxesMeta 

17 class Step(metaclass=AxesMeta): 

18 pass 

19 

20 # All subclasses automatically get axes support (metaclass inherited) 

21 class MyStep(Step, axes={"scope": "/pipeline/step_0"}): 

22 pass 

23 

24 MyStep.__axes__["scope"] # "/pipeline/step_0" 

25 

26 # Per-key MRO inheritance works automatically 

27 class ChildStep(MyStep, axes={"priority": 1}): 

28 pass 

29 

30 ChildStep.__axes__["scope"] # "/pipeline/step_0" (inherited) 

31 ChildStep.__axes__["priority"] # 1 (defined here) 

32""" 

33 

34from typing import Dict, Any, Tuple, Optional, Callable 

35from types import MappingProxyType 

36import weakref 

37 

38# ============================================================================= 

39# TYPE CACHE - Same pattern as reified_generics 

40# ============================================================================= 

41 

42_axes_cache: Dict[Tuple[type, tuple, tuple], type] = {} 

43 

44 

45def _cache_key(origin: type, bases: tuple, axes: dict) -> tuple: 

46 """Create hashable cache key from axes dict.""" 

47 # Sort axes for consistent hashing 

48 axes_tuple = tuple(sorted(axes.items())) 

49 return (origin, bases, axes_tuple) 

50 

51 

52# ============================================================================= 

53# AXES METACLASS 

54# ============================================================================= 

55 

56class AxesMeta(type): 

57 """ 

58 Metaclass prototype for type() with axes support. 

59 

60 In the real CPython implementation, this logic would be in type.__new__, 

61 and every class would have __axes__ = MappingProxyType({}) by default. 

62 

63 This metaclass is the stand-in: once any class uses metaclass=AxesMeta, 

64 all subclasses inherit it automatically (standard metaclass inheritance). 

65 """ 

66 

67 def __new__(mcs, name: str, bases: tuple, namespace: dict, 

68 axes: Optional[Dict[str, Any]] = None, **kwargs): 

69 """Create type with axes metadata and per-key MRO inheritance.""" 

70 # Create the class 

71 cls = super().__new__(mcs, name, bases, namespace) 

72 

73 # Per-key MRO inheritance: collect from parents, first wins 

74 parent_axes: Dict[str, Any] = {} 

75 for parent in cls.__mro__[1:]: 

76 for k, v in getattr(parent, '__axes__', {}).items(): 

77 if k not in parent_axes: 

78 parent_axes[k] = v 

79 

80 # Child axes override parent axes 

81 if axes: 

82 parent_axes.update(axes) 

83 

84 # Store as immutable mapping 

85 cls.__axes__ = MappingProxyType(parent_axes) 

86 

87 # Convenience attributes: __scope__, __registry__, etc. 

88 for axis_name, axis_value in parent_axes.items(): 

89 setattr(cls, f"__{axis_name}__", axis_value) 

90 

91 return cls 

92 

93 def __repr__(cls) -> str: 

94 if cls.__axes__: 

95 axes_str = ', '.join(f"{k}={v!r}" for k, v in cls.__axes__.items()) 

96 return f"<class '{cls.__name__}' axes=({axes_str})>" 

97 return super().__repr__() 

98 

99 

100# ============================================================================= 

101# DECORATOR + BASE CLASS (used by tests / ergonomic API) 

102# ============================================================================= 

103 

104 

105def with_axes(**axes: Any) -> Callable[[type], type]: 

106 """Class decorator that attaches axes metadata. 

107 

108 The decorated class is recreated using :class:`AxesMeta` so it gains: 

109 - ``__axes__`` (MappingProxyType) 

110 - convenience attributes (``__scope__``, ``__priority__``, etc.) 

111 """ 

112 

113 def _decorator(cls: type) -> type: 

114 # Recreate class using AxesMeta. Filter out internal descriptors. 

115 namespace = { 

116 k: v 

117 for k, v in cls.__dict__.items() 

118 if k not in {"__dict__", "__weakref__"} 

119 } 

120 return AxesMeta(cls.__name__, cls.__bases__, namespace, axes=axes) 

121 

122 return _decorator 

123 

124 

125class AxesBase(metaclass=AxesMeta): 

126 """Opt-in base class enabling ``class Foo(AxesBase, axes={...})`` syntax.""" 

127 

128 pass 

129 

130 

131# ============================================================================= 

132# FACTORY FUNCTION - Mimics extended type() signature 

133# ============================================================================= 

134 

135def axes_type(name: str, bases: tuple, namespace: dict, **axes) -> type: 

136 """ 

137 Create a type with arbitrary axes - prototype for extended type(). 

138  

139 This mimics the proposed signature: 

140 type(name, bases, namespace, **axes) 

141  

142 Args: 

143 name: Class name 

144 bases: Base classes tuple 

145 namespace: Class namespace dict 

146 **axes: Arbitrary axes (scope=, registry=, version=, etc.) 

147  

148 Returns: 

149 New type with __axes__ containing all axis values 

150  

151 Example: 

152 MyStep = axes_type("MyStep", (Step,), {"process": fn}, 

153 scope="/pipeline/step_0", 

154 registry=step_registry) 

155  

156 MyStep.__axes__["scope"] # "/pipeline/step_0" 

157 MyStep.__scope__ # "/pipeline/step_0" (convenience) 

158 """ 

159 return AxesMeta(name, bases, namespace, axes=axes) 

160 

161 

162# ============================================================================= 

163# UTILITY FUNCTIONS 

164# ============================================================================= 

165 

166def get_axes(cls: type) -> Dict[str, Any]: 

167 """Get all axes from a type.""" 

168 return getattr(cls, '__axes__', {}) 

169 

170 

171def get_axis(cls: type, name: str, default: Any = None) -> Any: 

172 """Get a specific axis value.""" 

173 return get_axes(cls).get(name, default) 

174 

175 

176def has_axis(cls: type, name: str) -> bool: 

177 """Check if type has a specific axis.""" 

178 return name in get_axes(cls) 

179