fix: resolve enum generation issues and test failures

Author: mikewolfd

pyproject.toml

@@ -1,10 +1,10 @@
 [build-system]
-requires = ["setuptools>=61.0", "wheel"]
-build-backend = "setuptools.build_meta"
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "msgspec-schemaorg"
-version = "0.2.0"
+version = "0.2.1"
 description = "Generate Python msgspec.Struct classes from the Schema.org vocabulary"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -36,6 +36,29 @@ dev = [
     "pytest"
 ]
 
+[tool.poetry]
+name = "msgspec-schemaorg"
+description = "Python msgspec.Struct models for Schema.org"
+authors = ["Mike Wolfson <mike@mikewolfson.com>"]
+readme = "README.md"
+homepage = "https://github.com/mikewolfd/msgspec-schemaorg"
+repository = "https://github.com/mikewolfd/msgspec-schemaorg"
+version = "0.2.1"
+packages = [{include = "msgspec_schemaorg"}]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+msgspec = ">=0.16.0"
+pydantic = ">=2.0.0"
+requests = ">=2.28.2"
+
+[tool.poetry.group.dev.dependencies]
+black = "^24.1.1"
+isort = "^5.13.2"
+pytest = "^8.0.1"
+mypy = "^1.8.0"
+ruff = "^0"
+
 [tool.setuptools]
 packages = ["msgspec_schemaorg", "msgspec_schemaorg.enums", "msgspec_schemaorg.enums.intangible"]
 include-package-data = true 

scripts/generate_enums.py

@@ -97,19 +97,18 @@ def generate_enum_code(enum_type: str, values: list, category: str = None):
         if ':' in value_id:
             value_id = value_id.split(':')[-1]
 
-        # Clean comment for JSON embedding
+        # Get comment and ensure it's a string
         comment = value.get('comment', '')
         if not isinstance(comment, str):
             comment = str(comment)
 
-        # Escape quotes in comment
-        comment = comment.replace('"', '\\"')
-        
+        # Use triple quotes for comment to properly handle multi-line text
         code.append(f'        "{value_id}": {{')
         code.append(f'            "id": "{value["id"]}",')
-        code.append(f'            "comment": "{comment}",')
+        code.append(f'            "comment": """{comment}""",')
         if value.get('label'):
-            code.append(f'            "label": "{value["label"]}",')
+            label = str(value["label"]).replace('"', '\\"')
+            code.append(f'            "label": "{label}",')
         code.append('        },')
 
     code.append("    }")

scripts/generate_models.py

@@ -68,6 +68,31 @@ def get_existing_enum_types():
     return enum_types
 
 
+def get_enum_categories():
+    """
+    Get a mapping of enum types to their categories.
+    
+    Returns:
+        A dictionary mapping enum type names to their categories
+    """
+    enum_categories = {}
+    
+    if not ENUMS_DIR.exists():
+        return enum_categories
+    
+    for root, _, files in os.walk(ENUMS_DIR):
+        category = os.path.basename(root)
+        if category == "enums":  # Skip the root enums directory
+            continue
+            
+        for file in files:
+            if file.endswith('.py') and not file.startswith('__'):
+                enum_name = os.path.splitext(file)[0]
+                enum_categories[enum_name] = category
+    
+    return enum_categories
+
+
 def modify_imports_for_enums(files, enum_types):
     """
     Modify import statements in generated files to use enum types from the enums package.
@@ -76,17 +101,8 @@ def modify_imports_for_enums(files, enum_types):
         files: Dictionary mapping file paths to generated code
         enum_types: Set of type names that have enum implementations
     """
-    enum_locations = {}
-    
-    # First, find all enum implementations and their locations
-    for root, _, files_in_dir in os.walk(ENUMS_DIR):
-        for file in files_in_dir:
-            if file.endswith('.py') and not file.startswith('__'):
-                enum_name = os.path.splitext(file)[0]
-                # Get the category from the directory structure
-                rel_path = os.path.relpath(root, ENUMS_DIR)
-                category = rel_path if rel_path != '.' else ''
-                enum_locations[enum_name] = category
+    # Get enum categories
+    enum_categories = get_enum_categories()
 
     # Now update imports in all files
     for file_path, content in files.items():
@@ -97,20 +113,21 @@ def modify_imports_for_enums(files, enum_types):
 
         # Look for import statements for enum types
         for enum_type in enum_types:
-            if enum_type not in enum_locations:
+            if enum_type not in enum_categories:
                 continue
 
-            category = enum_locations[enum_type]
-            category_path = f".{category}" if category else ""
+            category = enum_categories[enum_type]
 
             # Patterns for imports from models (handle various potential patterns)
             old_import_patterns = [
                 f"from msgspec_schemaorg.models.intangible.{enum_type} import {enum_type}",
-                f"from msgspec_schemaorg.models{category_path}.{enum_type} import {enum_type}"
+                f"from msgspec_schemaorg.models.{category}.{enum_type} import {enum_type}",
+                f"from .{enum_type} import {enum_type}",
+                f"from ..{category}.{enum_type} import {enum_type}"
             ]
 
             # Pattern for corrected import from enums
-            new_import = f"from msgspec_schemaorg.enums{category_path}.{enum_type} import {enum_type}"
+            new_import = f"from msgspec_schemaorg.enums.{category}.{enum_type} import {enum_type}"
 
             # Replace the import statement
             for old_pattern in old_import_patterns:
@@ -135,73 +152,28 @@ def modify_init_files(files, enum_types):
             continue
 
         modified_content = content
+        modified = False
 
         # Remove imports for enum types
         for enum_type in enum_types:
             import_line = f"from .{enum_type} import {enum_type}\n"
             if import_line in modified_content:
                 modified_content = modified_content.replace(import_line, "")
+                modified = True
                 print(f"Removed import for {enum_type} in {file_path}")
 
         # Update __all__ list if present
         if "__all__ = [" in modified_content:
             # Find __all__ list
             all_list_start = modified_content.find("__all__ = [")
             all_list_end = modified_content.find("]", all_list_start)
-            all_list = modified_content[all_list_start:all_list_end+1]
-            
-            # Remove enum types from __all__ list
-            for enum_type in enum_types:
-                # Look for different patterns in __all__ list
-                patterns = [
-                    f"'{enum_type}',",
-                    f"'{enum_type}'",
-                    f"\"{enum_type}\",",
-                    f"\"{enum_type}\""
-                ]
-                
-                for pattern in patterns:
-                    if pattern in all_list:
-                        # Replace with empty string or just a space depending on location
-                        if pattern.endswith(","):
-                            all_list = all_list.replace(pattern, "")
-                        else:
-                            all_list = all_list.replace(pattern, "")
-                        print(f"Removed {enum_type} from __all__ list in {file_path}")
-            
-            # Clean up any empty commas or double commas
-            all_list = all_list.replace(",,", ",")
-            all_list = all_list.replace(", ,", ",")
-            all_list = all_list.replace("[,", "[")
-            all_list = all_list.replace(",]", "]")
-            
-            # Replace the old __all__ list with the cleaned up one
-            modified_content = modified_content.replace(
-                modified_content[all_list_start:all_list_end+1],
-                all_list
-            )
-        
-        # Update the content in the files dictionary
-        files[file_path] = modified_content
-    
-    # Also update the root models/__init__.py if it exists
-    root_init_path = Path(DEFAULT_OUTPUT_DIR) / "__init__.py"
-    if root_init_path.exists():
-        try:
-            with open(root_init_path, 'r') as f:
-                root_init_content = f.read()
-            
-            modified_root_init = root_init_content
-            
-            # Update __all__ list if present
-            if "__all__ = [" in modified_root_init:
-                # Find __all__ list
-                all_list_start = modified_root_init.find("__all__ = [")
-                all_list_end = modified_root_init.find("]", all_list_start)
-                all_list = modified_root_init[all_list_start:all_list_end+1]
+            if all_list_end > all_list_start:  # Ensure we found the end bracket
+                all_list = modified_content[all_list_start:all_list_end+1]
+                original_all_list = all_list
 
                 # Remove enum types from __all__ list
                 for enum_type in enum_types:
+                    # Look for different patterns in __all__ list
                     patterns = [
                         f"'{enum_type}',",
                         f"'{enum_type}'",
@@ -213,7 +185,8 @@ def modify_init_files(files, enum_types):
                         if pattern in all_list:
                             # Replace with empty string
                             all_list = all_list.replace(pattern, "")
-                            print(f"Removed {enum_type} from __all__ list in root __init__.py")
+                            modified = True
+                            print(f"Removed {enum_type} from __all__ list in {file_path}")
 
                 # Clean up any empty commas or double commas
                 all_list = all_list.replace(",,", ",")
@@ -222,14 +195,68 @@ def modify_init_files(files, enum_types):
                 all_list = all_list.replace(",]", "]")
 
                 # Replace the old __all__ list with the cleaned up one
-                modified_root_init = modified_root_init.replace(
-                    modified_root_init[all_list_start:all_list_end+1],
-                    all_list
-                )
+                if all_list != original_all_list:
+                    modified_content = modified_content.replace(
+                        original_all_list,
+                        all_list
+                    )
+        
+        # Update the content in the files dictionary if modified
+        if modified:
+            files[file_path] = modified_content
+    
+    # Also update the root models/__init__.py if it exists
+    root_init_path = Path(DEFAULT_OUTPUT_DIR) / "__init__.py"
+    if root_init_path.exists():
+        try:
+            with open(root_init_path, 'r') as f:
+                root_init_content = f.read()
+            
+            modified_root_init = root_init_content
+            modified = False
+            
+            # Update __all__ list if present
+            if "__all__ = [" in modified_root_init:
+                # Find __all__ list
+                all_list_start = modified_root_init.find("__all__ = [")
+                all_list_end = modified_root_init.find("]", all_list_start)
+                if all_list_end > all_list_start:  # Ensure we found the end bracket
+                    all_list = modified_root_init[all_list_start:all_list_end+1]
+                    original_all_list = all_list
+                    
+                    # Remove enum types from __all__ list
+                    for enum_type in enum_types:
+                        patterns = [
+                            f"'{enum_type}',",
+                            f"'{enum_type}'",
+                            f"\"{enum_type}\",",
+                            f"\"{enum_type}\""
+                        ]
+                        
+                        for pattern in patterns:
+                            if pattern in all_list:
+                                # Replace with empty string
+                                all_list = all_list.replace(pattern, "")
+                                modified = True
+                                print(f"Removed {enum_type} from __all__ list in root __init__.py")
+                    
+                    # Clean up any empty commas or double commas
+                    all_list = all_list.replace(",,", ",")
+                    all_list = all_list.replace(", ,", ",")
+                    all_list = all_list.replace("[,", "[")
+                    all_list = all_list.replace(",]", "]")
+                    
+                    # Replace the old __all__ list with the cleaned up one
+                    if all_list != original_all_list:
+                        modified_root_init = modified_root_init.replace(
+                            original_all_list,
+                            all_list
+                        )
 
-                # Write the modified content back to the file
-                with open(root_init_path, 'w') as f:
-                    f.write(modified_root_init)
+                # Write the modified content back to the file if changed
+                if modified:
+                    with open(root_init_path, 'w') as f:
+                        f.write(modified_root_init)
         except Exception as e:
             print(f"Error updating root __init__.py: {e}")
 

tests/test_schema_examples.py

@@ -4,13 +4,15 @@
 """
 
 import importlib
+import importlib.util
+import pkgutil  # Import pkgutil for iterating modules
 import inspect
 import json
 import os
 import re
 import sys
 from collections.abc import Iterator
-from typing import Union, Any, Dict, List, Optional
+from typing import Union, Any, Dict, List, Optional, Set, Type, get_type_hints
 
 import requests
 from requests import Response
@@ -54,22 +56,31 @@ def get_modules_in_package(package_name: str) -> Iterator[Any]:
     # Get all modules in the package
     if hasattr(package, "__path__"):
         package_path = package.__path__
-        for _, module_name, is_pkg in importlib.iter_modules(package_path):
-            full_module_name = f"{package_name}.{module_name}"
-            try:
-                module = importlib.import_module(full_module_name)
-                
-                # If this is a package, get all classes from it
-                if is_pkg:
-                    yield from get_modules_in_package(full_module_name)
-                
-                # Get all classes from this module
-                for name in dir(module):
-                    attr = getattr(module, name)
-                    if isinstance(attr, type) and issubclass(attr, msgspec.Struct) and attr.__module__ == module.__name__:
-                        yield attr
-            except ImportError:
-                continue
+        for module_info in pkgutil.iter_modules(package_path):  # Use pkgutil.iter_modules
+            module_name = module_info.name  # Access name attribute
+            is_pkg = module_info.ispkg     # Access ispkg attribute
+            
+            if is_pkg:
+                # If module is a package, recurse into it
+                sub_package = f"{package_name}.{module_name}"
+                try:
+                    yield from get_modules_in_package(sub_package)
+                except (ImportError, ModuleNotFoundError):
+                    continue
+            else:
+                try:
+                    # Try to import the module
+                    module = importlib.import_module(f"{package_name}.{module_name}")
+                    
+                    # Find all class objects defined in the module
+                    for name, obj in inspect.getmembers(module):
+                        if inspect.isclass(obj) and obj.__module__ == module.__name__:
+                            # Skip internal classes
+                            if not name.startswith("_"):
+                                yield obj
+                except (ImportError, ModuleNotFoundError):
+                    # Skip if module can't be imported
+                    continue
     else:
         # This is a module, not a package
         for name in dir(package):