wishful 0.2.1__py3-none-any.whl

wishful/types/registry.py

@@ -0,0 +1,333 @@
+"""Type registration system for wishful.
+
+Allows users to register complex types (Pydantic models, dataclasses, TypedDict)
+that the LLM can use when generating code.
+"""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import fields as dataclass_fields
+from dataclasses import is_dataclass
+from typing import Any, Callable, TypeVar, get_type_hints
+
+T = TypeVar("T")
+
+
+class TypeRegistry:
+    """Global registry for user-defined types."""
+
+    def __init__(self):
+        # Map: type_name -> serialized type definition
+        self._types: dict[str, str] = {}
+        # Map: function_name -> type_name (for output_for mapping)
+        self._function_outputs: dict[str, str] = {}
+
+    def register(
+        self, type_class: type, *, output_for: str | list[str] | None = None
+    ) -> None:
+        """Register a type and optionally associate it with function(s)."""
+        schema = self._serialize_type(type_class)
+        self._types[type_class.__name__] = schema
+
+        if output_for:
+            functions = [output_for] if isinstance(output_for, str) else output_for
+            for func_name in functions:
+                self._function_outputs[func_name] = type_class.__name__
+
+    def get_schema(self, type_name: str) -> str | None:
+        """Get the serialized schema for a registered type."""
+        return self._types.get(type_name)
+
+    def get_all_schemas(self) -> dict[str, str]:
+        """Get all registered type schemas."""
+        return self._types.copy()
+
+    def get_output_type(self, function_name: str) -> str | None:
+        """Get the registered output type for a function."""
+        return self._function_outputs.get(function_name)
+
+    def clear(self) -> None:
+        """Clear all registered types."""
+        self._types.clear()
+        self._function_outputs.clear()
+
+    def _serialize_type(self, type_class: type) -> str:
+        """Serialize a type to a string representation for the LLM."""
+        # Check if it's a Pydantic model
+        if self._is_pydantic_model(type_class):
+            return self._serialize_pydantic(type_class)
+
+        # Check if it's a dataclass
+        if is_dataclass(type_class):
+            return self._serialize_dataclass(type_class)
+
+        # Check if it's a TypedDict
+        if self._is_typed_dict(type_class):
+            return self._serialize_typed_dict(type_class)
+
+        # Fallback: get source code if available
+        try:
+            return inspect.getsource(type_class)
+        except (OSError, TypeError):
+            # Last resort: just return the class definition line
+            return f"class {type_class.__name__}: ..."
+
+    def _is_pydantic_model(self, type_class: type) -> bool:
+        """Check if a class is a Pydantic BaseModel."""
+        try:
+            # Check if BaseModel is in the MRO or has model_fields
+            return hasattr(type_class, "model_fields") or any(
+                "BaseModel" in base.__name__ for base in type_class.__mro__
+            )
+        except (AttributeError, TypeError):
+            return False
+
+    def _serialize_pydantic(self, model_class: type) -> str:
+        """Serialize a Pydantic model to source code."""
+        lines = [f"class {model_class.__name__}(BaseModel):"]
+
+        # Add docstring if present
+        if model_class.__doc__:
+            lines.append(f'    """{model_class.__doc__.strip()}"""')
+
+        # Get model fields
+        if hasattr(model_class, "model_fields"):
+            # Pydantic v2
+            for field_name, field_info in model_class.model_fields.items():
+                annotation = self._format_annotation(field_info.annotation)
+
+                # Check if field is required (has no default)
+                # Use callable check for robustness with mocks
+                is_required = field_info.is_required() if callable(getattr(field_info, 'is_required', None)) else (field_info.default is None and field_info.default_factory is None)
+                
+                # Check if field has metadata (Field() usage)
+                has_field_metadata = hasattr(field_info, 'metadata') and field_info.metadata
+                has_constraints = any(
+                    hasattr(field_info, attr) and getattr(field_info, attr) is not None
+                    for attr in ['description', 'min_length', 'max_length', 'gt', 'ge', 'lt', 'le', 'pattern']
+                )
+                
+                if is_required:
+                    if has_field_metadata or has_constraints:
+                        # Build Field() arguments
+                        field_args = self._build_field_args(field_info)
+                        lines.append(f"    {field_name}: {annotation} = Field({field_args})")
+                    else:
+                        lines.append(f"    {field_name}: {annotation}")
+                else:
+                    # Field has a default value or default_factory
+                    if field_info.default_factory is not None:
+                        field_args = self._build_field_args(field_info)
+                        if field_args:
+                            lines.append(
+                                f"    {field_name}: {annotation} = Field(default_factory=..., {field_args})"
+                            )
+                        else:
+                            lines.append(
+                                f"    {field_name}: {annotation} = Field(default_factory=...)"
+                            )
+                    else:
+                        # Has a default value
+                        if has_field_metadata or has_constraints:
+                            field_args = self._build_field_args(field_info)
+                            default_repr = repr(field_info.default)
+                            lines.append(f"    {field_name}: {annotation} = Field(default={default_repr}, {field_args})")
+                        else:
+                            default_repr = repr(field_info.default)
+                            lines.append(f"    {field_name}: {annotation} = {default_repr}")
+        elif hasattr(model_class, "__fields__"):
+            # Pydantic v1
+            for field_name, field in model_class.__fields__.items():
+                annotation = self._format_annotation(field.outer_type_)
+                if field.required:
+                    lines.append(f"    {field_name}: {annotation}")
+                else:
+                    default_repr = repr(field.default)
+                    lines.append(f"    {field_name}: {annotation} = {default_repr}")
+
+        return "\n".join(lines)
+    
+    def _build_field_args(self, field_info) -> str:
+        """Build Field() arguments from field_info metadata and constraints."""
+        args = []
+        
+        # Add description
+        if hasattr(field_info, 'description') and field_info.description:
+            args.append(f"description={repr(field_info.description)}")
+        
+        # Parse metadata for Pydantic v2 constraints
+        if hasattr(field_info, 'metadata') and field_info.metadata:
+            for meta_item in field_info.metadata:
+                meta_class = meta_item.__class__.__name__
+                
+                # Handle common constraint types
+                if meta_class == 'MinLen' and hasattr(meta_item, 'min_length'):
+                    args.append(f"min_length={meta_item.min_length}")
+                elif meta_class == 'MaxLen' and hasattr(meta_item, 'max_length'):
+                    args.append(f"max_length={meta_item.max_length}")
+                elif meta_class == 'Gt' and hasattr(meta_item, 'gt'):
+                    args.append(f"gt={meta_item.gt}")
+                elif meta_class == 'Ge' and hasattr(meta_item, 'ge'):
+                    args.append(f"ge={meta_item.ge}")
+                elif meta_class == 'Lt' and hasattr(meta_item, 'lt'):
+                    args.append(f"lt={meta_item.lt}")
+                elif meta_class == 'Le' and hasattr(meta_item, 'le'):
+                    args.append(f"le={meta_item.le}")
+                elif meta_class == '_PydanticGeneralMetadata':
+                    # Handle pattern and other general metadata
+                    if hasattr(meta_item, 'pattern') and meta_item.pattern:
+                        args.append(f"pattern={repr(meta_item.pattern)}")
+        
+        # Fallback: check direct attributes (Pydantic v1 or custom)
+        constraints = [
+            ('min_length', 'min_length'),
+            ('max_length', 'max_length'),
+            ('gt', 'gt'),
+            ('ge', 'ge'),
+            ('lt', 'lt'),
+            ('le', 'le'),
+            ('pattern', 'pattern'),
+        ]
+        
+        for attr_name, arg_name in constraints:
+            if hasattr(field_info, attr_name):
+                value = getattr(field_info, attr_name)
+                if value is not None and arg_name not in ' '.join(args):  # Avoid duplicates
+                    args.append(f"{arg_name}={repr(value)}")
+        
+        return ", ".join(args)
+
+    def _serialize_dataclass(self, dc_class: type) -> str:
+        """Serialize a dataclass to source code."""
+        lines = ["@dataclass", f"class {dc_class.__name__}:"]
+
+        # Add docstring if present
+        if dc_class.__doc__:
+            lines.append(f'    """{dc_class.__doc__.strip()}"""')
+
+        for field in dataclass_fields(dc_class):
+            annotation = self._format_annotation(field.type)
+            if field.default is not field.default_factory:  # type: ignore
+                # Has a default value
+                default_repr = repr(field.default)
+                lines.append(f"    {field.name}: {annotation} = {default_repr}")
+            elif field.default_factory is not field.default_factory:  # type: ignore
+                lines.append(
+                    f"    {field.name}: {annotation} = field(default_factory=...)"
+                )
+            else:
+                lines.append(f"    {field.name}: {annotation}")
+
+        return "\n".join(lines)
+
+    def _is_typed_dict(self, type_class: type) -> bool:
+        """Check if a class is a TypedDict."""
+        try:
+            return hasattr(type_class, "__annotations__") and hasattr(
+                type_class, "__total__"
+            )
+        except AttributeError:
+            return False
+
+    def _serialize_typed_dict(self, td_class: type) -> str:
+        """Serialize a TypedDict to source code."""
+        lines = [f"class {td_class.__name__}(TypedDict):"]
+
+        if td_class.__doc__:
+            lines.append(f'    """{td_class.__doc__.strip()}"""')
+
+        for field_name, field_type in get_type_hints(td_class).items():
+            annotation = self._format_annotation(field_type)
+            lines.append(f"    {field_name}: {annotation}")
+
+        return "\n".join(lines)
+
+    def _format_annotation(self, annotation: Any) -> str:
+        """Format a type annotation as a string."""
+        if hasattr(annotation, "__name__"):
+            return annotation.__name__
+
+        # Handle typing generics
+        if hasattr(annotation, "__origin__"):
+            origin = annotation.__origin__
+            args = getattr(annotation, "__args__", ())
+
+            if origin is list:
+                return (
+                    f"list[{self._format_annotation(args[0])}]" if args else "list"
+                )
+            elif origin is dict:
+                key_type = self._format_annotation(args[0]) if args else "Any"
+                val_type = self._format_annotation(args[1]) if len(args) > 1 else "Any"
+                return f"dict[{key_type}, {val_type}]"
+            elif origin is tuple:
+                arg_strs = ", ".join(self._format_annotation(a) for a in args)
+                return f"tuple[{arg_strs}]"
+            # Handle Union/Optional
+            elif hasattr(origin, "__name__") and origin.__name__ == "UnionType":
+                arg_strs = " | ".join(self._format_annotation(a) for a in args)
+                return arg_strs
+
+        return str(annotation).replace("typing.", "")
+
+
+# Global registry instance
+_registry = TypeRegistry()
+
+
+def type(
+    cls: type[T] | None = None, *, output_for: str | list[str] | None = None
+) -> type[T] | Callable[[type[T]], type[T]]:
+    """Decorator to register a type with wishful.
+
+    Usage:
+        @wishful.type
+        class UserProfile(BaseModel):
+            name: str
+            email: str
+
+        # Or with output type specification
+        @wishful.type(output_for='create_user')
+        class UserProfile(BaseModel):
+            name: str
+            email: str
+
+        # Multiple functions
+        @wishful.type(output_for=['create_user', 'update_user'])
+        class UserProfile(BaseModel):
+            name: str
+            email: str
+    """
+
+    def decorator(type_class: type[T]) -> type[T]:
+        _registry.register(type_class, output_for=output_for)
+        return type_class
+
+    # Handle both @wishful.type and @wishful.type(...) syntax
+    if cls is None:
+        # Called with arguments: @wishful.type(output_for='...')
+        return decorator
+    else:
+        # Called without arguments: @wishful.type
+        return decorator(cls)
+
+
+def get_type_schema(type_name: str) -> str | None:
+    """Get the schema for a registered type."""
+    return _registry.get_schema(type_name)
+
+
+def get_all_type_schemas() -> dict[str, str]:
+    """Get all registered type schemas."""
+    return _registry.get_all_schemas()
+
+
+def get_output_type_for_function(function_name: str) -> str | None:
+    """Get the output type registered for a function."""
+    return _registry.get_output_type(function_name)
+
+
+def clear_type_registry() -> None:
+    """Clear all registered types (useful for testing)."""
+    _registry.clear()

wishful/ui.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from wishful.config import settings
+
+_console = Console()
+
+
+@contextmanager
+def spinner(message: str) -> Iterator[None]:
+    if not settings.spinner:
+        yield
+        return
+
+    with Progress(SpinnerColumn(), TextColumn(message), console=_console, transient=True) as progress:
+        task_id = progress.add_task(message, total=None)
+        try:
+            yield
+        finally:
+            progress.update(task_id, completed=1)
+            progress.stop()