feat: add create_grammar_model()

craft_grammar/__init__.py

@@ -27,6 +27,7 @@
 from ._statement import CallStack, Grammar, Statement
 from ._to import ToStatement
 from ._try import TryStatement
+from .create import create_grammar_model
 
 __all__ = [
     "errors",
@@ -38,4 +39,5 @@
     "Statement",
     "ToStatement",
     "TryStatement",
+    "create_grammar_model",
 ]

craft_grammar/create.py

@@ -0,0 +1,152 @@
+#  This file is part of craft-grammar.
+#
+#  Copyright 2024 Canonical Ltd.
+#
+#  This program is free software: you can redistribute it and/or modify it
+#  under the terms of the GNU Lesser General Public License version 3, as
+#  published by the Free Software Foundation.
+#
+#  This program is distributed in the hope that it will be useful, but WITHOUT
+#  ANY WARRANTY; without even the implied warranties of MERCHANTABILITY,
+#  SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR PURPOSE.
+#  See the GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+"""Utilities to create grammar models."""
+import builtins
+import logging
+import types
+import typing
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+CONFIG_TEMPLATE = """
+    class Config:
+        validate_assignment = True
+        extra = "ignore"
+        allow_mutation = False
+        alias_generator = lambda s: s.replace("_", "-")
+"""
+
+
+def create_grammar_model(model_class: typing.Type[BaseModel]) -> str:
+    """Create the code for a grammar-aware class compatible with ``model_class``.
+
+    :param model_class: A pydantic.BaseModel subclass.
+    """
+    class_decl = f"class Grammar{model_class.__name__}(BaseModel):"
+
+    attributes = []
+    type_hints = typing.get_type_hints(model_class)
+    for attr_name, attr_type in type_hints.items():
+        if attr_name == "__slots__":
+            # This happens in Python 3.12 apparently
+            continue
+
+        grammar_type = _get_grammar_type_for(attr_type)
+
+        if grammar_type is None:
+            logger.debug(
+                "Skipping unknown type %s for attribute %s", attr_type, attr_name
+            )
+            continue
+
+        attr_field = model_class.__fields__[attr_name]
+        alias = attr_field.alias
+        new_name = alias if "-" not in alias else attr_name
+        attr_decl = f"{new_name}: {grammar_type}"
+
+        if not attr_field.required:
+            default_factory = attr_field.default_factory
+            if default_factory is not None:
+                default = repr(default_factory())
+            else:
+                default = repr(attr_field.default)
+            # repr(x) uses single quotes for strings; replace them with double
+            # quotes here to be consistent with the codebase.
+            default = default.replace("'", '"')
+            attr_decl += f" = {default}"
+
+        attributes.append(attr_decl)
+
+    lines = [class_decl] + CONFIG_TEMPLATE.split("\n")
+    for attr in attributes:
+        lines.append(f"    {attr}")
+    lines.append("")  # Final newline
+
+    return "\n".join(lines)
+
+
+def _get_grammar_type_for(model_type) -> str | None:  # pylint: disable=R0911,R0912
+    """Get the "grammar" type for ``model_type``.
+
+    Returns None if we don't know how to "grammify" ``model_type``.
+    """
+    if model_type is type(None):
+        # None -> None
+        return "None"
+
+    origin = typing.get_origin(model_type)
+    args = typing.get_args(model_type)
+
+    match origin:
+        case None:
+            # Primitive, regular class, Pydantic model, etc.
+            if issubclass(model_type, BaseModel):
+                # PydanticModel -> GrammarPydanticModel
+                # (assumes that generate_grammar_model(model_type) will be called).
+                return f"Grammar{model_type.__name__}"
+            return f"Grammar[{model_type.__name__}]"
+
+        case typing.Union:
+            # Type is either a Union[] or an Optional[]
+            if len(args) == 2 and type(None) in args:
+                # Type is an Optional[]
+                # Optional[T] -> Optional[Grammar[T]]
+                other_type = [t for t in args if t is not type(None)][0]
+                grammar_type = _get_grammar_type_for(other_type)
+                return f"Optional[{grammar_type}]"
+
+            # Union[X, Y] -> Grammar[Union[X, Y]]
+            union_args = []
+            for arg in typing.get_args(model_type):
+                if typing.get_origin(arg) is None:
+                    if arg is type(None):
+                        name = "None"
+                    else:
+                        # print int as "int"
+                        name = arg.__name__
+                    union_args.append(name)
+                else:
+                    # print dict[k, v] as "dict[k,v]"
+                    union_args.append(str(arg))
+            comma_args = ", ".join(union_args)
+            return f"Grammar[Union[{comma_args}]]"
+
+        case types.UnionType:
+            # Type is an expression like "str | int | None"
+            # A type like "str | None" becomes "Grammar[str] | None"
+            grammar_types = [_get_grammar_type_for(a) for a in args]
+            grammar_args = [t for t in grammar_types if t is not None]
+            return " | ".join(grammar_args)
+
+        case builtins.list | builtins.dict:
+            # list[T] -> Grammar[list[T]]
+            name = str(model_type).removeprefix("typing.")
+            return f"Grammar[{name}]"
+
+        case typing.Literal:
+            # Literal["a", "b"] -> Grammar[str]
+            arg_types = set(type(a) for a in typing.get_args(model_type))
+            if len(arg_types) == 1:
+                # For now only handle the case where all possible literal values
+                # have the same type
+                arg_type = arg_types.pop()
+                return _get_grammar_type_for(arg_type)
+            return None
+
+        case _:
+            return None

setup.cfg

@@ -29,6 +29,7 @@ packages = find:
 zip_safe = False
 install_requires =
     overrides
+    pydantic<2.0
 
 [options.package_data]
 craft_grammar = py.typed
@@ -51,7 +52,6 @@ test =
     flake8
     isort==5.10.1
     mypy==0.991
-    pydantic<2.0
     pydocstyle==6.1.1
     pylint==2.15.10
     pylint-fixme-info

tests/unit/test_create.py

@@ -0,0 +1,90 @@
+#  This file is part of craft-grammar.
+#
+#  Copyright 2024 Canonical Ltd.
+#
+#  This program is free software: you can redistribute it and/or modify it
+#  under the terms of the GNU Lesser General Public License version 3, as
+#  published by the Free Software Foundation.
+#
+#  This program is distributed in the hope that it will be useful, but WITHOUT
+#  ANY WARRANTY; without even the implied warranties of MERCHANTABILITY,
+#  SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR PURPOSE.
+#  See the GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+from __future__ import annotations
+
+from typing import List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from craft_grammar import create_grammar_model
+
+
+class SubModel(BaseModel):
+    """A Pydantic model used as an attribute of another model."""
+
+
+class MyModel(BaseModel):
+    """A Pydantic model that we want to "grammify"."""
+
+    # Primitive types
+    str_value: str
+    str_value_or_none: str | None
+    optional_str_value: Optional[str]
+    str_with_default: str = "string"
+    str_or_non_with_default: str | None = "string or None"
+    union_value: Union[str, int, None]
+    literal_value: Literal["red", "green", "blue"] = "green"
+
+    # Collections
+    list_value: list[int] = []
+    other_list: List[int]
+    dict_value: dict[str, bool]
+    list_of_dicts: list[dict[str, str]]
+
+    # Pydantic models
+    sub_model: SubModel
+
+    # An aliased field (should use the alias)
+    aliased_field: int = Field(default=1, alias="alias_name")
+
+    # A field with a default factory
+    factory_field: str = Field(default_factory=str)
+
+    # Untyped fields are ignored
+    untyped_1 = "untyped_1"
+    untyped_2 = 1
+
+
+EXPECTED_GRAMMAR_MODEL = """\
+class GrammarMyModel(BaseModel):
+
+    class Config:
+        validate_assignment = True
+        extra = "ignore"
+        allow_mutation = False
+        alias_generator = lambda s: s.replace("_", "-")
+
+    str_value: Grammar[str]
+    str_value_or_none: Grammar[str] | None = None
+    optional_str_value: Optional[Grammar[str]] = None
+    str_with_default: Grammar[str] = "string"
+    str_or_non_with_default: Grammar[str] | None = "string or None"
+    union_value: Grammar[Union[str, int, None]] = None
+    literal_value: Grammar[str] = "green"
+    list_value: Grammar[list[int]] = []
+    other_list: Grammar[List[int]]
+    dict_value: Grammar[dict[str, bool]]
+    list_of_dicts: Grammar[list[dict[str, str]]]
+    sub_model: GrammarSubModel
+    alias_name: Grammar[int] = 1
+    factory_field: Grammar[str] = ""
+"""
+
+
+def test_create_model():
+    grammar_model = create_grammar_model(MyModel)
+
+    assert grammar_model == EXPECTED_GRAMMAR_MODEL