Coverage for src/invariant/registry.py: 85.94%

1"""OpRegistry for mapping operation names to callables."""

2

3 import types

4 from importlib.metadata import entry_points

5 from typing import Any, Callable

6

7# Type alias for op packages: dict mapping short names to op callables

8 OpPackage = dict[str, Callable[..., Any]]

9

10

11 class OpRegistry:

12 """Singleton registry mapping string identifiers to executable Python callables.

13

14 Decouples the "string" name in the graph definition from the actual Python code.

15 """

16

17 _instance: "OpRegistry | None" = None

18 _initialized: bool = False

19

20 def __new__(cls) -> "OpRegistry":

21 """Ensure singleton pattern."""

22 if cls._instance is None:

23 cls._instance = super().__new__(cls)

24 return cls._instance

25

26 def __init__(self) -> None:

27 """Initialize the registry (only once)."""

28 if not OpRegistry._initialized:

29 self._ops: dict[str, Callable[..., Any]] = {}

30 OpRegistry._initialized = True

31

32 def register(self, name: str, op: Callable[..., Any]) -> None:

33 """Register an operation.

34

35 Args:

36 name: The string identifier for the operation.

37 op: The callable that implements the operation.

38 Should be a plain Python function with typed parameters.

39

40 Raises:

41 ValueError: If name is empty or already registered.

42 """

43 if not name:

44 raise ValueError("Operation name cannot be empty")

45 if name in self._ops:

46 raise ValueError(f"Operation '{name}' is already registered")

47 self._ops[name] = op

48

49 def get(self, name: str) -> Callable[..., Any]:

50 """Get an operation by name.

51

52 Args:

53 name: The string identifier for the operation.

54

55 Returns:

56 The callable that implements the operation.

57

58 Raises:

59 KeyError: If operation is not registered.

60 """

61 if name not in self._ops:

62 raise KeyError(f"Operation '{name}' is not registered")

63 return self._ops[name]

64

65 def has(self, name: str) -> bool:

66 """Check if an operation is registered.

67

68 Args:

69 name: The string identifier for the operation.

70

71 Returns:

72 True if registered, False otherwise.

73 """

74 return name in self._ops

75

76 def clear(self) -> None:

77 """Clear all registered operations (mainly for testing)."""

78 self._ops.clear()

79

80 def register_package(self, prefix: str, ops: OpPackage | Any) -> None:

81 """Register all ops from a package under a common prefix.

82

83 Args:

84 prefix: The namespace prefix (e.g. "poly").

85 ops: Either a dict mapping short names to callables (OpPackage),

86 or a Python module that has an OPS dict attribute.

87

88 Raises:

89 ValueError: If prefix is empty, ops is invalid, or any operation

90 name is already registered.

91 AttributeError: If ops is a module but doesn't have an OPS attribute.

92 """

93 if not prefix:

94 raise ValueError("Package prefix cannot be empty")

95

96 # Extract the ops dict from the input

97 ops_dict: OpPackage

98 if isinstance(ops, dict):

99 ops_dict = ops

100 elif isinstance(ops, types.ModuleType):

101 # It's a module - check for OPS attribute

102 if not hasattr(ops, "OPS"):

103 raise AttributeError(

104 f"Module {ops.__name__} does not have an OPS attribute"

105 )

106 ops_dict = ops.OPS

107 if not isinstance(ops_dict, dict):

108 raise ValueError(f"OPS attribute must be a dict, got {type(ops_dict)}")

109 elif hasattr(ops, "OPS"):

110 # Object with OPS attribute (not a module)

111 ops_dict = ops.OPS

112 if not isinstance(ops_dict, dict):

113 raise ValueError(f"OPS attribute must be a dict, got {type(ops_dict)}")

114 else:

115 raise ValueError(

116 f"ops must be a dict or module with OPS attribute, got {type(ops)}"

117 )

118

119 # Register each op with the prefix

120 for name, op in ops_dict.items():

121 full_name = f"{prefix}:{name}"

122 self.register(full_name, op)

123

124 def auto_discover(self) -> None:

125 """Discover and register op packages from entry points.

126

127 Scans the 'invariant.ops' entry point group. Each entry point

128 should resolve to either:

129 - A dict[str, Callable] (the OPS dict directly)

130 - A callable that returns such a dict

131

132 The entry point name becomes the package prefix.

133

134 Raises:

135 ValueError: If any operation name is already registered (via register_package).

136 """

137 eps = entry_points(group="invariant.ops")

138

139 for ep in eps:

140 try:

141 # Load the entry point

142 loaded = ep.load()

143

144 # Extract the ops dict

145 ops_dict: OpPackage

146 if isinstance(loaded, dict):

147 ops_dict = loaded

148 elif callable(loaded):

149 # Callable that returns the dict

150 result = loaded()

151 if not isinstance(result, dict):

152 continue # Skip invalid entry points

153 ops_dict = result

154 else:

155 continue # Skip invalid entry points

156

157 # Register the package using the entry point name as prefix

158 self.register_package(ep.name, ops_dict)

159 except Exception:

160 # Skip invalid entry points silently

161 continue