Flowfile 0.3.10__py3-none-any.whl → 0.4.0__py3-none-any.whl

flowfile_core/configs/node_store/user_defined_node_registry.py

@@ -0,0 +1,193 @@
+import sys
+import importlib.util
+import inspect
+from pathlib import Path
+from typing import Dict, Type, List
+import logging
+
+from flowfile_core.flowfile.node_designer.custom_node import CustomNodeBase, NodeSettings
+from shared import storage
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_all_custom_nodes() -> Dict[str, Type[CustomNodeBase]]:
+    """
+    Scan the user-defined nodes directory and import all CustomNodeBase subclasses.
+
+    Returns:
+        Dictionary mapping node names to node classes
+    """
+    custom_nodes = {}
+
+    # Get the directory path where user-defined nodes are stored
+    nodes_directory = storage.user_defined_nodes_icons
+
+    # Convert to Path object for easier handling
+    nodes_path = Path(nodes_directory)
+
+    if not nodes_path.exists() or not nodes_path.is_dir():
+        print(f"Warning: Nodes directory {nodes_path} does not exist or is not a directory")
+        return custom_nodes
+
+    # Scan all Python files in the directory
+    for file_path in nodes_path.glob("*.py"):
+        # Skip __init__.py and other special files
+        if file_path.name.startswith("__"):
+            continue
+
+        try:
+            # Load the module dynamically
+            module_name = file_path.stem  # filename without extension
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+
+            if spec and spec.loader:
+                module = importlib.util.module_from_spec(spec)
+
+                # Add to sys.modules to handle imports within the module
+                sys.modules[module_name] = module
+
+                # Execute the module
+                spec.loader.exec_module(module)
+
+                # Inspect the module for CustomNodeBase subclasses
+                for name, obj in inspect.getmembers(module):
+                    # Check if it's a class and a subclass of CustomNodeBase
+                    # but not CustomNodeBase itself
+                    if (inspect.isclass(obj) and
+                            issubclass(obj, CustomNodeBase) and
+                            obj is not CustomNodeBase):
+
+                        # Use the node_name attribute if it exists, otherwise use class name
+                        node_name = getattr(obj, 'node_name', name)
+                        custom_nodes[node_name] = obj
+                        print(f"Loaded custom node: {node_name} from {file_path.name}")
+
+        except Exception as e:
+            print(f"Error loading module from {file_path}: {e}")
+            # Continue with other files even if one fails
+            continue
+
+    return custom_nodes
+
+
+def get_all_custom_nodes_with_validation() -> Dict[str, Type[CustomNodeBase]]:
+    """
+    Enhanced version that validates the nodes before adding them.
+    """
+
+    custom_nodes = {}
+    nodes_path = storage.user_defined_nodes_directory
+
+    if not nodes_path.exists():
+        return custom_nodes
+
+    for file_path in nodes_path.glob("*.py"):
+        if file_path.name.startswith("__"):
+            continue
+
+        try:
+            module_name = file_path.stem
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+
+            if spec and spec.loader:
+                module = importlib.util.module_from_spec(spec)
+                sys.modules[module_name] = module
+                spec.loader.exec_module(module)
+
+                for name, obj in inspect.getmembers(module):
+                    if (inspect.isclass(obj) and
+                            issubclass(obj, CustomNodeBase) and
+                            obj is not CustomNodeBase):
+
+                        try:
+                            _obj = obj()
+                            # Validate that the node has required attributes
+                            if not hasattr(_obj, 'node_name'):
+                                logger.error(f"Warning: {name} missing node_name attribute")
+                                raise ValueError(f"Node {name} must implement a node_name attribute")
+
+                            if not hasattr(_obj, 'settings_schema'):
+                                logger.error(f"Warning: {name} missing settings_schema attribute")
+                                raise ValueError(f"Node {name} must implement a settings_schema attribute")
+
+                            if not hasattr(_obj, 'process'):
+                                logger.error(f"Warning: {name} missing process method")
+                                raise ValueError(f"Node {name} must implement a process method")
+                            if not (storage.user_defined_nodes_icons / _obj.node_icon).exists():
+                                logger.warning(
+                                    f"Warning: Icon file does not exist for node {_obj.node_name} at {_obj.node_icon} "
+                                    "Falling back to default icon."
+                                )
+
+                            node_name = _obj.to_node_template().item
+                            custom_nodes[node_name] = obj
+                            print(f"✓ Loaded: {node_name} from {file_path.name}")
+                        except Exception as e:
+                            print(f"Error validating node {name} in {file_path}: {e}")
+                            continue
+        except SyntaxError as e:
+            print(f"Syntax error in {file_path}: {e}")
+        except ImportError as e:
+            print(f"Import error in {file_path}: {e}")
+        except Exception as e:
+            print(f"Unexpected error loading {file_path}: {e}")
+
+    return custom_nodes
+
+
+def get_custom_nodes_lazy() -> List[Type[CustomNodeBase]]:
+    """
+    Returns a list of custom node classes without instantiating them.
+    Useful for registration or catalog purposes.
+    """
+    nodes = []
+    nodes_path = Path(storage.user_defined_nodes_directory)
+
+    if not nodes_path.exists():
+        return nodes
+
+    for file_path in nodes_path.glob("*.py"):
+        if file_path.name.startswith("__"):
+            continue
+
+        try:
+            # Create a unique module name to avoid conflicts
+            module_name = f"custom_node_{file_path.stem}_{id(file_path)}"
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+
+            if spec and spec.loader:
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+
+                for name, obj in inspect.getmembers(module):
+                    if (inspect.isclass(obj) and
+                            issubclass(obj, CustomNodeBase) and
+                            obj is not CustomNodeBase and
+                            obj.__module__ == module.__name__):  # Only get classes defined in this module
+                        nodes.append(obj)
+
+        except Exception as e:
+            print(f"Error processing {file_path}: {e}")
+            continue
+
+    return nodes
+
+
+# Example usage function that matches your original pattern
+def add_custom_node(node_class: Type[CustomNodeBase], registry: Dict[str, Type[CustomNodeBase]]):
+    """Add a single custom node to the registry."""
+    if hasattr(node_class, 'node_name'):
+        registry[node_class.node_name] = node_class
+    else:
+        registry[node_class.__name__] = node_class
+
+
+def get_all_nodes_from_standard_location() -> Dict[str, Type[CustomNodeBase]]:
+    """
+    Main function to get all custom nodes from the standard location.
+    This matches your original function signature.
+    """
+
+    return get_all_custom_nodes_with_validation()

flowfile_core/flowfile/flow_data_engine/flow_file_column/interface.py

@@ -0,0 +1,4 @@
+from typing import Literal
+
+DataTypeGroup = Literal['numeric', 'str', 'date']
+ReadableDataTypeGroup = Literal['Numeric', 'String', 'Date', 'Other', 'Boolean', 'Binary', 'Complex']

flowfile_core/flowfile/flow_data_engine/flow_file_column/main.py

@@ -1,44 +1,13 @@
 
 from dataclasses import dataclass
-from typing import Optional, Any, List, Dict, Literal, Iterable
+from typing import Optional, Any, List, Dict, Iterable
 
 from flowfile_core.schemas import input_schema
 from flowfile_core.flowfile.flow_data_engine.flow_file_column.utils import cast_str_to_polars_type
 from flowfile_core.flowfile.flow_data_engine.flow_file_column.polars_type import PlType
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.interface import ReadableDataTypeGroup, DataTypeGroup
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.type_registry import convert_pl_type_to_string
 import polars as pl
-# TODO: rename flow_file_column to flowfile_column
-DataTypeGroup = Literal['numeric', 'str', 'date']
-
-
-def convert_pl_type_to_string(pl_type: pl.DataType, inner: bool = False) -> str:
-    if isinstance(pl_type, pl.List):
-        inner_str = convert_pl_type_to_string(pl_type.inner, inner=True)
-        return f"pl.List({inner_str})"
-    elif isinstance(pl_type, pl.Array):
-        inner_str = convert_pl_type_to_string(pl_type.inner, inner=True)
-        return f"pl.Array({inner_str})"
-    elif isinstance(pl_type, pl.Decimal):
-        precision = pl_type.precision if hasattr(pl_type, 'precision') else None
-        scale = pl_type.scale if hasattr(pl_type, 'scale') else None
-        if precision is not None and scale is not None:
-            return f"pl.Decimal({precision}, {scale})"
-        elif precision is not None:
-            return f"pl.Decimal({precision})"
-        else:
-            return "pl.Decimal()"
-    elif isinstance(pl_type, pl.Struct):
-        # Handle Struct with field definitions
-        fields = []
-        if hasattr(pl_type, 'fields'):
-            for field in pl_type.fields:
-                field_name = field.name
-                field_type = convert_pl_type_to_string(field.dtype, inner=True)
-                fields.append(f'pl.Field("{field_name}", {field_type})')
-        field_str = ", ".join(fields)
-        return f"pl.Struct([{field_str}])"
-    else:
-        # For base types, we want the full pl.TypeName format
-        return str(pl_type.base_type()) if not inner else f"pl.{pl_type}"
 
 
 @dataclass
@@ -52,6 +21,7 @@ class FlowfileColumn:
     number_of_empty_values: int
     number_of_unique_values: int
     example_values: str
+    data_type_group: ReadableDataTypeGroup
     __sql_type: Optional[Any]
     __is_unique: Optional[bool]
     __nullable: Optional[bool]
@@ -75,6 +45,7 @@ class FlowfileColumn:
         self.__is_unique = None
         self.__sql_type = None
         self.__perc_unique = None
+        self.data_type_group = self.get_readable_datatype_group()
 
     def __repr__(self):
         """
@@ -220,6 +191,20 @@ class FlowfileColumn:
             return 'numeric'
         elif self.data_type in ('datetime', 'date', 'Date', 'Datetime', 'Time'):
             return 'date'
+        else:
+            return 'str'
+
+    def get_readable_datatype_group(self) -> ReadableDataTypeGroup:
+        if self.data_type in ('Utf8', 'VARCHAR', 'CHAR', 'NVARCHAR', 'String'):
+            return 'String'
+        elif self.data_type in ('fixed_decimal', 'decimal', 'float', 'integer', 'boolean', 'double', 'Int16', 'Int32',
+                                'Int64', 'Float32', 'Float64', 'Decimal', 'Binary', 'Boolean', 'Uint8', 'Uint16',
+                                'Uint32', 'Uint64'):
+            return 'Numeric'
+        elif self.data_type in ('datetime', 'date', 'Date', 'Datetime', 'Time'):
+            return 'Date'
+        else:
+            return 'Other'
 
     def get_polars_type(self) -> PlType:
         pl_datatype = cast_str_to_polars_type(self.data_type)

flowfile_core/flowfile/flow_data_engine/flow_file_column/type_registry.py

@@ -0,0 +1,36 @@
+
+from typing import Type, Literal, List, Dict, Union, Tuple
+import polars as pl
+DataTypeGroup = Literal['numeric', 'string', 'datetime', 'boolean', 'binary', 'complex', 'unknown']
+
+
+def convert_pl_type_to_string(pl_type: pl.DataType, inner: bool = False) -> str:
+    if isinstance(pl_type, pl.List):
+        inner_str = convert_pl_type_to_string(pl_type.inner, inner=True)
+        return f"pl.List({inner_str})"
+    elif isinstance(pl_type, pl.Array):
+        inner_str = convert_pl_type_to_string(pl_type.inner, inner=True)
+        return f"pl.Array({inner_str})"
+    elif isinstance(pl_type, pl.Decimal):
+        precision = pl_type.precision if hasattr(pl_type, 'precision') else None
+        scale = pl_type.scale if hasattr(pl_type, 'scale') else None
+        if precision is not None and scale is not None:
+            return f"pl.Decimal({precision}, {scale})"
+        elif precision is not None:
+            return f"pl.Decimal({precision})"
+        else:
+            return "pl.Decimal()"
+    elif isinstance(pl_type, pl.Struct):
+        # Handle Struct with field definitions
+        fields = []
+        if hasattr(pl_type, 'fields'):
+            for field in pl_type.fields:
+                field_name = field.name
+                field_type = convert_pl_type_to_string(field.dtype, inner=True)
+                fields.append(f'pl.Field("{field_name}", {field_type})')
+        field_str = ", ".join(fields)
+        return f"pl.Struct([{field_str}])"
+    else:
+        # For base types, we want the full pl.TypeName format
+        return str(pl_type.base_type()) if not inner else f"pl.{pl_type}"
+

flowfile_core/flowfile/flow_graph.py

@@ -16,7 +16,7 @@ from pyarrow.parquet import ParquetFile
 from flowfile_core.configs import logger
 from flowfile_core.configs.flow_logger import FlowLogger
 from flowfile_core.flowfile.sources.external_sources.factory import data_source_factory
-from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import cast_str_to_polars_type, FlowfileColumn
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn, cast_str_to_polars_type
 
 from flowfile_core.flowfile.flow_data_engine.cloud_storage_reader import CloudStorageReader
 from flowfile_core.utils.arrow_reader import get_read_top_n
@@ -51,6 +51,7 @@ from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source impor
 from flowfile_core.flowfile.database_connection_manager.db_connections import (get_local_database_connection,
                                                                                get_local_cloud_connection)
 from flowfile_core.flowfile.util.calculate_layout import calculate_layered_layout
+from flowfile_core.flowfile.node_designer.custom_node import CustomNodeBase
 
 
 def get_xlsx_schema(engine: str, file_path: str, sheet_name: str, start_row: int, start_column: int,
@@ -436,6 +437,24 @@ class FlowGraph:
         node = self._node_db.get(node_id)
         if node is not None:
             return node
+        
+    def add_user_defined_node(self, *,
+                              custom_node: CustomNodeBase,
+                              user_defined_node_settings: input_schema.UserDefinedNode
+                              ):
+       
+        def _func(*fdes: FlowDataEngine) -> FlowDataEngine | None:
+            output = custom_node.process(*(fde.data_frame for fde in fdes))
+            if isinstance(output, pl.LazyFrame | pl.DataFrame):
+                return FlowDataEngine(output)
+            return None
+        
+        self.add_node_step(node_id=user_defined_node_settings.node_id,
+                           function=_func,
+                           setting_input=user_defined_node_settings,
+                           input_node_ids=user_defined_node_settings.depending_on_ids,
+                           node_type=custom_node.item,
+                           )
 
     def add_pivot(self, pivot_settings: input_schema.NodePivot):
         """Adds a pivot node to the graph.

flowfile_core/flowfile/flow_node/flow_node.py

@@ -8,7 +8,7 @@ from flowfile_core.configs.flow_logger import NodeLogger
 
 from flowfile_core.schemas.output_model import TableExample, FileColumn, NodeData
 from flowfile_core.flowfile.utils import get_hash
-from flowfile_core.configs.node_store import nodes as node_interface
+from flowfile_core.configs import node_store
 from flowfile_core.flowfile.setting_generator import setting_generator, setting_updator
 from time import sleep
 from flowfile_core.flowfile.flow_data_engine.subprocess_operations import (
@@ -27,7 +27,7 @@ class FlowNode:
     """
     parent_uuid: str
     node_type: str
-    node_template: node_interface.NodeTemplate
+    node_template: node_store.NodeTemplate
     node_default: schemas.NodeDefault
     node_schema: NodeSchemaInformation
     node_inputs: NodeStepInputs
@@ -251,10 +251,10 @@ class FlowNode:
         self.results.errors = None
         self.add_lead_to_in_depend_source()
         _ = self.hash
-        self.node_template = node_interface.node_dict.get(self.node_type)
+        self.node_template = node_store.node_dict.get(self.node_type)
         if self.node_template is None:
             raise Exception(f'Node template {self.node_type} not found')
-        self.node_default = node_interface.node_defaults.get(self.node_type)
+        self.node_default = node_store.node_defaults.get(self.node_type)
         self.setting_input = setting_input  # wait until the end so that the hash is calculated correctly
 
     @property

flowfile_core/flowfile/manage/open_flowfile.py

@@ -4,6 +4,7 @@ from flowfile_core.flowfile.manage.compatibility_enhancements import ensure_comp
 import pickle
 from flowfile_core.flowfile.flow_graph import FlowGraph
 from pathlib import Path
+from flowfile_core.configs.node_store import CUSTOM_NODE_STORE
 
 
 def determine_insertion_order(node_storage: schemas.FlowInformation):
@@ -81,7 +82,14 @@ def open_flow(flow_path: Path) -> FlowGraph:
         new_flow.add_node_promise(node_promise)
     for node_id in ingestion_order:
         node_info: schemas.NodeInformation = flow_storage_obj.data[node_id]
-        getattr(new_flow, 'add_' + node_info.type)(node_info.setting_input)
+        if hasattr(node_info.setting_input, "is_user_defined") and node_info.setting_input.is_user_defined:
+            if node_info.type not in CUSTOM_NODE_STORE:
+                continue
+            user_defined_node_class = CUSTOM_NODE_STORE[node_info.type]
+            new_flow.add_user_defined_node(custom_node=user_defined_node_class.from_settings(node_info.setting_input.settings),
+                                           user_defined_node_settings=node_info.setting_input)
+        else:
+            getattr(new_flow, 'add_' + node_info.type)(node_info.setting_input)
         from_node = new_flow.get_node(node_id)
         for output_node_id in node_info.outputs:
             to_node = new_flow.get_node(output_node_id)

flowfile_core/flowfile/node_designer/__init__.py

@@ -0,0 +1,47 @@
+# flowfile_core/flowfile/node_designer/__init__.py
+
+"""
+Tools for creating custom Flowfile nodes.
+
+This package provides all the necessary components for developers to build their own
+custom nodes, define their UI, and implement their data processing logic.
+"""
+
+# Import the core base class for creating a new node
+from .custom_node import CustomNodeBase, NodeSettings
+
+# Import all UI components so they can be used directly
+from .ui_components import (
+    Section,
+    TextInput,
+    NumericInput,
+    ToggleSwitch,
+    SingleSelect,
+    MultiSelect,
+    ColumnSelector,
+    IncomingColumns,  # Important marker class for dynamic dropdowns
+)
+
+# Import the main `Types` object for filtering in ColumnSelector
+from .data_types import Types
+
+
+# Define the public API of this package
+__all__ = [
+    # Core Node Class
+    "CustomNodeBase",
+
+    # UI Components & Layout
+    "Section",
+    "TextInput",
+    "NumericInput",
+    "ToggleSwitch",
+    "SingleSelect",
+    "MultiSelect",
+    "NodeSettings",
+    "ColumnSelector",
+    "IncomingColumns",
+
+    # Data Type Filtering
+    "Types",
+]

flowfile_core/flowfile/node_designer/_type_registry.py

@@ -0,0 +1,197 @@
+# _type_registry.py - Internal type system (not for public use)
+"""
+Internal type registry for mapping between different type representations.
+This module should not be imported directly by users.
+"""
+
+from dataclasses import dataclass
+from typing import Type, List, Dict, Set, Any, Union
+import polars as pl
+
+# Import public types
+from flowfile_core.flowfile.node_designer.data_types import TypeGroup, DataType
+
+
+@dataclass(frozen=True)
+class TypeMapping:
+    """Internal mapping between type representations."""
+    data_type: DataType
+    polars_type: Type[pl.DataType]
+    type_group: TypeGroup
+    aliases: tuple[str, ...] = ()
+
+
+class TypeRegistry:
+    """
+    Internal registry for type conversions and lookups.
+    This class is not part of the public API.
+    """
+
+    def __init__(self):
+        self._mappings: List[TypeMapping] = [
+            # Numeric types
+            TypeMapping(DataType.Int8, pl.Int8, TypeGroup.Numeric, ("i8",)),
+            TypeMapping(DataType.Int16, pl.Int16, TypeGroup.Numeric, ("i16",)),
+            TypeMapping(DataType.Int32, pl.Int32, TypeGroup.Numeric, ("i32", "int32")),
+            TypeMapping(DataType.Int64, pl.Int64, TypeGroup.Numeric,
+                        ("i64", "int64", "int", "integer", "bigint")),
+            TypeMapping(DataType.UInt8, pl.UInt8, TypeGroup.Numeric, ("u8",)),
+            TypeMapping(DataType.UInt16, pl.UInt16, TypeGroup.Numeric, ("u16",)),
+            TypeMapping(DataType.UInt32, pl.UInt32, TypeGroup.Numeric, ("u32", "uint32")),
+            TypeMapping(DataType.UInt64, pl.UInt64, TypeGroup.Numeric, ("u64", "uint64")),
+            TypeMapping(DataType.Float32, pl.Float32, TypeGroup.Numeric, ("f32", "float32")),
+            TypeMapping(DataType.Float64, pl.Float64, TypeGroup.Numeric,
+                        ("f64", "float64", "float", "double")),
+            TypeMapping(DataType.Decimal, pl.Decimal, TypeGroup.Numeric,
+                        ("decimal", "numeric", "dec")),
+
+            # String types
+            TypeMapping(DataType.String, pl.String, TypeGroup.String,
+                        ("str", "string", "utf8", "varchar", "text")),
+            TypeMapping(DataType.Categorical, pl.Categorical, TypeGroup.String,
+                        ("cat", "categorical", "enum", "factor")),
+
+            # Date types
+            TypeMapping(DataType.Date, pl.Date, TypeGroup.Date, ("date",)),
+            TypeMapping(DataType.Datetime, pl.Datetime, TypeGroup.Date,
+                        ("datetime", "timestamp")),
+            TypeMapping(DataType.Time, pl.Time, TypeGroup.Date, ("time",)),
+            TypeMapping(DataType.Duration, pl.Duration, TypeGroup.Date,
+                        ("duration", "timedelta")),
+
+            # Other types
+            TypeMapping(DataType.Boolean, pl.Boolean, TypeGroup.Boolean,
+                        ("bool", "boolean")),
+            TypeMapping(DataType.Binary, pl.Binary, TypeGroup.Binary,
+                        ("binary", "bytes", "bytea")),
+            TypeMapping(DataType.List, pl.List, TypeGroup.Complex, ("list", "array")),
+            TypeMapping(DataType.Struct, pl.Struct, TypeGroup.Complex, ("struct", "object")),
+            TypeMapping(DataType.Array, pl.Array, TypeGroup.Complex, ("fixed_array",)),
+        ]
+
+        self._build_indices()
+
+    def _build_indices(self):
+        """Build lookup indices for fast access."""
+        self._by_data_type: Dict[DataType, TypeMapping] = {}
+        self._by_polars_type: Dict[Type[pl.DataType], TypeMapping] = {}
+        self._by_alias: Dict[str, TypeMapping] = {}
+        self._by_group: Dict[TypeGroup, List[TypeMapping]] = {g: [] for g in TypeGroup}
+
+        for mapping in self._mappings:
+            self._by_data_type[mapping.data_type] = mapping
+            self._by_polars_type[mapping.polars_type] = mapping
+
+            if mapping.type_group != TypeGroup.All:
+                self._by_group[mapping.type_group].append(mapping)
+
+            # Register all aliases (case-insensitive)
+            for alias in mapping.aliases:
+                self._by_alias[alias.lower()] = mapping
+
+            # Register enum names as aliases
+            self._by_alias[mapping.data_type.value.lower()] = mapping
+            self._by_alias[mapping.polars_type.__name__.lower()] = mapping
+
+            # Register "pl.TypeName" format
+            self._by_alias[f"pl.{mapping.polars_type.__name__}".lower()] = mapping
+
+    def normalize(self, type_spec: Any) -> Set[DataType]:
+        """
+        Normalize any type specification to a set of DataType enums.
+        This is the main internal API for type resolution.
+        """
+        # Handle special case: All types
+        if type_spec == TypeGroup.All or type_spec == "ALL":
+            return set(self._by_data_type.keys())
+
+        # Handle TypeGroup
+        if isinstance(type_spec, TypeGroup):
+            return {m.data_type for m in self._by_group.get(type_spec, [])}
+
+        # Handle DataType
+        if isinstance(type_spec, DataType):
+            return {type_spec}
+
+        # Handle Polars type class
+        if isinstance(type_spec, type) and issubclass(type_spec, pl.DataType):
+            mapping = self._by_polars_type.get(type_spec)
+            if mapping:
+                return {mapping.data_type}
+
+        # Handle Polars type instance
+        if isinstance(type_spec, pl.DataType):
+            base_type = type_spec.base_type() if hasattr(type_spec, 'base_type') else type(type_spec)
+            mapping = self._by_polars_type.get(base_type)
+            if mapping:
+                return {mapping.data_type}
+
+        # Handle string aliases
+        if isinstance(type_spec, str):
+            type_spec_lower = type_spec.lower()
+
+            # Try TypeGroup name
+            try:
+                group = TypeGroup(type_spec)
+                return {m.data_type for m in self._by_group.get(group, [])}
+            except (ValueError, KeyError):
+                pass
+
+            # Try DataType name
+            try:
+                dt = DataType(type_spec)
+                return {dt}
+            except (ValueError, KeyError):
+                pass
+
+            # Check aliases
+            mapping = self._by_alias.get(type_spec_lower)
+            if mapping:
+                return {mapping.data_type}
+
+        # Default to empty set if unrecognized
+        return set()
+
+    def normalize_list(self, type_specs: List[Any]) -> Set[DataType]:
+        """Normalize a list of type specifications."""
+        result = set()
+        for spec in type_specs:
+            result.update(self.normalize(spec))
+        return result
+
+    def get_polars_types(self, data_types: Set[DataType]) -> Set[Type[pl.DataType]]:
+        """Convert a set of DataType enums to Polars types."""
+        result = set()
+        for dt in data_types:
+            mapping = self._by_data_type.get(dt)
+            if mapping:
+                result.add(mapping.polars_type)
+        return result
+
+    def get_polars_type(self, data_type: DataType) -> Type[pl.DataType]:
+        """Get the Polars type for a single DataType."""
+        mapping = self._by_data_type.get(data_type)
+        return mapping.polars_type if mapping else pl.String  # Default fallback
+
+
+# Singleton instance
+_registry = TypeRegistry()
+
+
+# Internal API functions (not for public use)
+def normalize_type_spec(type_spec: Any) -> Set[DataType]:
+    """Internal function to normalize type specifications."""
+    if isinstance(type_spec, list):
+        return _registry.normalize_list(type_spec)
+    return _registry.normalize(type_spec)
+
+
+def get_polars_types(data_types: Set[DataType]) -> Set[Type[pl.DataType]]:
+    """Internal function to get Polars types."""
+    return _registry.get_polars_types(data_types)
+
+
+def check_column_type(column_dtype: pl.DataType, accepted_types: Set[DataType]) -> bool:
+    """Check if a column's dtype matches the accepted types."""
+    normalized = _registry.normalize(column_dtype)
+    return bool(normalized & accepted_types)