Component Registry

Register and resolve components by name.

See Also

• Config Overview - All config tools
• Loaders - Config file loaders
• Schema - Config validation
• Core Config - Config protocol

---

datarax.config.registry

Component registry for Datarax.

This module provides a registry system for Datarax components, allowing components to be registered and instantiated dynamically from configuration. Enhanced with Flax NNX module support for proper RNG and Variable handling.

logger module-attribute

logger = getLogger(__name__)

ComponentConstructor module-attribute

ComponentConstructor = Callable[..., Any]

register_component

register_component(component_type: str, name: str | None = None) -> Callable[[_T], _T]

Register a component constructor with the registry.

This decorator can be used to register a component constructor with the component registry, allowing it to be instantiated from configuration. Enhanced to handle both regular classes and NNX modules.

Parameters:

Name	Type	Description	Default
component_type	str	The type of component being registered (e.g., "source", "transformer", "augmenter").	required
name	str | None	Optional custom name for the component. If not provided, the class name will be used.	None

Returns:

Type	Description
Callable[[_T], _T]	A decorator function that registers the component constructor.

Examples:

@register_component("source")
class MyDataSource(DataSourceModule):
    def __init__(self, path: str, *, rngs: nnx.Rngs | None = None):
        super().__init__(rngs=rngs)
        self.path = path

get_component_constructor

get_component_constructor(component_type: str, name: str) -> ComponentConstructor

Get a component constructor by type and name.

Parameters:

Name	Type	Description	Default
component_type	str	The type of component (e.g., "source", "transformer").	required
name	str	The name of the component.	required

Returns:

Type	Description
ComponentConstructor	The component constructor function or class.

Raises:

Type	Description
KeyError	If the component is not registered.

is_component_registered

is_component_registered(component_type: str, name: str) -> bool

Check if a component is registered.

Parameters:

Name	Type	Description	Default
component_type	str	The type of component (e.g., "source", "transformer").	required
name	str	The name of the component.	required

Returns:

Type	Description
bool	True if the component is registered, False otherwise.

list_registered_components

list_registered_components(component_type: str | None = None) -> dict[str, ComponentConstructor]

List all registered components of a given type.

Parameters:

Name	Type	Description	Default
component_type	str | None	Optional type of components to list. If None, list all registered components.	None

Returns:

Type	Description
dict[str, ComponentConstructor]	Dictionary mapping component names to constructors.

create_component_from_config

create_component_from_config(component_type: str, name: str, config: dict[str, Any]) -> Any

Create a component instance from configuration.

Enhanced to properly handle NNX modules with RNG and Variable initialization.

Parameters:

Name	Type	Description	Default
component_type	str	The type of component (e.g., "source", "transformer").	required
name	str	The name of the component.	required
config	dict[str, Any]	Configuration dictionary for the component.	required

Returns:

Type	Description
Any	The instantiated component.

Raises:

Type	Description
KeyError	If the component is not registered.
TypeError	If the configuration is invalid.

get_component_info

get_component_info(component_type: str, name: str) -> dict[str, Any]

Get detailed information about a registered component.

Parameters:

Name	Type	Description	Default
component_type	str	The type of component.	required
name	str	The name of the component.	required

Returns:

Type	Description
dict[str, Any]	Dictionary containing component information including whether it's an NNX module,
dict[str, Any]	required parameters, and RNG requirements.