Data Types and Type System

class UInt64(Integer):
    """64-bit unsigned integer type."""

class UInt32(Integer):
    """32-bit unsigned integer type."""

class UInt16(Integer):
    """16-bit unsigned integer type (alias: USMALLINT)."""

class UInt8(Integer):
    """8-bit unsigned integer type."""

class UTinyInteger(Integer):
    """Tiny unsigned integer type (alias: UInt1)."""
    name = "UTinyInt"

class TinyInteger(Integer):
    """Tiny signed integer type (alias: Int1)."""
    name = "TinyInt"

class USmallInteger(Integer):
    """Unsigned small integer type (alias: UInt2)."""
    name = "usmallint"
    min = 0
    max = 65535

class UBigInteger(Integer):
    """Unsigned big integer type."""
    name = "UBigInt"
    min = 0
    max = 18446744073709551615

class HugeInteger(Integer):
    """128-bit signed integer type."""
    name = "HugeInt"

class UHugeInteger(Integer):
    """128-bit unsigned integer type."""
    name = "UHugeInt"

class UInteger(Integer):
    """Unsigned 32-bit integer type."""

class VarInt(Integer):
    """Variable-width integer type (DuckDB v1.0+)."""

class Struct(TypeEngine):
    """
    Represents a STRUCT type in DuckDB for nested data structures.
    
    Allows defining structured data with named fields of different types.
    """
    __visit_name__ = "struct"
    
    def __init__(self, fields=None):
        """
        Initialize STRUCT type.
        
        Parameters:
        - fields (Dict[str, Union[Type[TypeEngine], TypeEngine]], optional): 
          Dictionary mapping field names to SQLAlchemy types
        """
        self.fields = fields

class Map(TypeEngine):
    """
    Represents a MAP type in DuckDB for key-value pair data.
    
    Stores key-value mappings with specified key and value types.
    """
    __visit_name__ = "map"
    
    def __init__(self, key_type, value_type):
        """
        Initialize MAP type.
        
        Parameters:
        - key_type (Union[Type[TypeEngine], TypeEngine]): Type for map keys
        - value_type (Union[Type[TypeEngine], TypeEngine]): Type for map values
        """
        self.key_type = key_type
        self.value_type = value_type
    
    def bind_processor(self, dialect):
        """
        Process Python dict values for DuckDB MAP format.
        
        Returns:
        Callable: Function to convert dict to DuckDB MAP format
        """
    
    def result_processor(self, dialect, coltype):
        """
        Process DuckDB MAP results back to Python dict.
        
        Returns:
        Callable: Function to convert DuckDB MAP to dict
        """

class Union(TypeEngine):
    """
    Represents a UNION type in DuckDB for multiple possible types.
    
    Allows a column to contain values of different types.
    """
    __visit_name__ = "union"
    
    def __init__(self, fields):
        """
        Initialize UNION type.
        
        Parameters:
        - fields (Dict[str, Union[Type[TypeEngine], TypeEngine]]): 
          Dictionary mapping variant names to SQLAlchemy types
        """
        self.fields = fields

def register_extension_types():
    """
    Register all DuckDB extension types with SQLAlchemy compiler.
    
    This function sets up the necessary type compilation rules for
    DuckDB-specific types to work with SQLAlchemy's type system.
    """

def compile_uint(element, compiler, **kw):
    """
    Compile unsigned integer types to DuckDB SQL.
    
    Parameters:
    - element (Integer): Integer type instance
    - compiler (PGTypeCompiler): SQLAlchemy compiler
    - **kw: Additional compiler keyword arguments
    
    Returns:
    str: DuckDB type name
    """

ISCHEMA_NAMES: Dict[str, Type[TypeEngine]]
"""
Maps DuckDB type names to SQLAlchemy type classes.

Includes mappings for:
- Integer types: hugeint, uhugeint, tinyint, utinyint, etc.
- Floating point: float4, float8
- Temporal: timetz, timestamptz, timestamp_s, timestamp_ms, timestamp_ns
- Other: enum, bool, varchar
"""

types: List[Type[Integer]]
"""List of all custom integer type subclasses defined in this module."""

def process_type(value, compiler, **kw):
    """
    Process a type value through SQLAlchemy compiler.
    
    Parameters:
    - value (Union[TypeEngine, Type[TypeEngine]]): Type to process
    - compiler (PGTypeCompiler): SQLAlchemy compiler
    - **kw: Additional compiler arguments
    
    Returns:
    str: Compiled type string
    """

def struct_or_union(instance, compiler, identifier_preparer, **kw):
    """
    Generate SQL for STRUCT or UNION type definitions.
    
    Parameters:
    - instance (Union[Struct, Union]): Type instance
    - compiler (PGTypeCompiler): SQLAlchemy compiler  
    - identifier_preparer (PGIdentifierPreparer): Identifier preparer
    - **kw: Additional arguments
    
    Returns:
    str: Type definition SQL
    """

from duckdb_engine.datatypes import UInt64, HugeInteger, UTinyInteger
from sqlalchemy import Table, Column, MetaData

metadata = MetaData()
table = Table('data', metadata,
    Column('big_id', UInt64),
    Column('huge_number', HugeInteger), 
    Column('tiny_flag', UTinyInteger)
)

from duckdb_engine.datatypes import Struct
from sqlalchemy import Table, Column, String, Integer

# Define a struct with named fields
person_struct = Struct({
    'first_name': String(50),
    'last_name': String(50), 
    'age': Integer
})

table = Table('people', metadata,
    Column('id', Integer, primary_key=True),
    Column('person_info', person_struct)
)

from duckdb_engine.datatypes import Map
from sqlalchemy import String, Integer

# Define a map from string keys to integer values
tag_counts = Map(String, Integer)

table = Table('articles', metadata,
    Column('id', Integer, primary_key=True),
    Column('title', String(200)),
    Column('tag_counts', tag_counts)
)

from duckdb_engine.datatypes import Union
from sqlalchemy import String, Integer, Float

# Define a union that can hold string, int, or float values
mixed_value = Union({
    'text': String,
    'number': Integer,
    'decimal': Float
})

table = Table('flexible_data', metadata,
    Column('id', Integer, primary_key=True),
    Column('value', mixed_value)
)

from duckdb_engine.datatypes import register_extension_types

# This is automatically called when importing duckdb_engine
# but can be called manually if needed
register_extension_types()