Pydantic v1 and v2 support

.github/actions/poetry_setup/action.yml

@@ -0,0 +1,91 @@
+# An action for setting up poetry install with caching.
+# Using a custom action since the default action does not
+# take poetry install groups into account.
+# Action code from:
+# https://github.com/actions/setup-python/issues/505#issuecomment-1273013236
+name: poetry-install-with-caching
+description: Poetry install with support for caching of dependency groups.
+
+inputs:
+  python-version:
+    description: Python version, supporting MAJOR.MINOR only
+    required: true
+
+  poetry-version:
+    description: Poetry version
+    required: true
+
+  cache-key:
+    description: Cache key to use for manual handling of caching
+    required: true
+
+  working-directory:
+    description: Directory whose poetry.lock file should be cached
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/setup-python@v4
+      name: Setup python ${{ inputs.python-version }}
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - uses: actions/cache@v3
+      id: cache-bin-poetry
+      name: Cache Poetry binary - Python ${{ inputs.python-version }}
+      env:
+        SEGMENT_DOWNLOAD_TIMEOUT_MIN: "1"
+      with:
+        path: |
+          /opt/pipx/venvs/poetry
+        # This step caches the poetry installation, so make sure it's keyed on the poetry version as well.
+        key: bin-poetry-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-${{ inputs.poetry-version }}
+
+    - name: Refresh shell hashtable and fixup softlinks
+      if: steps.cache-bin-poetry.outputs.cache-hit == 'true'
+      shell: bash
+      env:
+        POETRY_VERSION: ${{ inputs.poetry-version }}
+        PYTHON_VERSION: ${{ inputs.python-version }}
+      run: |
+        set -eux
+
+        # Refresh the shell hashtable, to ensure correct `which` output.
+        hash -r
+
+        # `actions/cache@v3` doesn't always seem able to correctly unpack softlinks.
+        # Delete and recreate the softlinks pipx expects to have.
+        rm /opt/pipx/venvs/poetry/bin/python
+        cd /opt/pipx/venvs/poetry/bin
+        ln -s "$(which "python$PYTHON_VERSION")" python
+        chmod +x python
+        cd /opt/pipx_bin/
+        ln -s /opt/pipx/venvs/poetry/bin/poetry poetry
+        chmod +x poetry
+
+        # Ensure everything got set up correctly.
+        /opt/pipx/venvs/poetry/bin/python --version
+        /opt/pipx_bin/poetry --version
+
+    - name: Install poetry
+      if: steps.cache-bin-poetry.outputs.cache-hit != 'true'
+      shell: bash
+      env:
+        POETRY_VERSION: ${{ inputs.poetry-version }}
+        PYTHON_VERSION: ${{ inputs.python-version }}
+      run: pipx install "poetry==$POETRY_VERSION" --python "python$PYTHON_VERSION" --verbose
+
+    - name: Restore pip and poetry cached dependencies
+      uses: actions/cache@v3
+      env:
+        SEGMENT_DOWNLOAD_TIMEOUT_MIN: "4"
+        WORKDIR: ${{ inputs.working-directory == '' && '.' || inputs.working-directory }}
+      with:
+        path: |
+          ~/.cache/pip
+          ~/.cache/pypoetry/virtualenvs
+          ~/.cache/pypoetry/cache
+          ~/.cache/pypoetry/artifacts
+          ${{ env.WORKDIR }}/.venv
+        key: py-deps-${{ runner.os }}-${{ runner.arch }}-py-${{ inputs.python-version }}-poetry-${{ inputs.poetry-version }}-${{ inputs.cache-key }}-${{ hashFiles(format('{0}/**/poetry.lock', env.WORKDIR)) }}

.github/workflows/_pydantic_compatibility.yml

@@ -0,0 +1,81 @@
+name: pydantic v1/v2 compatibility
+
+on:
+  workflow_call:
+    inputs:
+      working-directory:
+        required: true
+        type: string
+        description: "From which folder this pipeline executes"
+
+env:
+  POETRY_VERSION: "1.5.1"
+
+jobs:
+  build:
+    defaults:
+      run:
+        working-directory: ${{ inputs.working-directory }}
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version:
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
+    name: Pydantic v1/v2 compatibility - Python ${{ matrix.python-version }}
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}
+        uses: "./.github/actions/poetry_setup"
+        with:
+          python-version: ${{ matrix.python-version }}
+          poetry-version: ${{ env.POETRY_VERSION }}
+          working-directory: ${{ inputs.working-directory }}
+          cache-key: pydantic-cross-compat
+
+      - name: Install dependencies
+        shell: bash
+        run: poetry install
+
+      - name: Install the opposite major version of pydantic
+        # If normal tests use pydantic v1, here we'll use v2, and vice versa.
+        shell: bash
+        run: |
+          # Determine the major part of pydantic version
+          REGULAR_VERSION=$(poetry run python -c "import pydantic; print(pydantic.__version__)" | cut -d. -f1)
+
+          if [[ "$REGULAR_VERSION" == "1" ]]; then
+            PYDANTIC_DEP=">=2.1,<3"
+            TEST_WITH_VERSION="2"
+          elif [[ "$REGULAR_VERSION" == "2" ]]; then
+            PYDANTIC_DEP="<2"
+            TEST_WITH_VERSION="1"
+          else
+            echo "Unexpected pydantic major version '$REGULAR_VERSION', cannot determine which version to use for cross-compatibility test."
+            exit 1
+          fi
+
+          # Install via `pip` instead of `poetry add` to avoid changing lockfile,
+          # which would prevent caching from working: the cache would get saved
+          # to a different key than where it gets loaded from.
+          poetry run pip install "pydantic${PYDANTIC_DEP}"
+
+          # Ensure that the correct pydantic is installed now.
+          echo "Checking pydantic version... Expecting ${TEST_WITH_VERSION}"
+
+          # Determine the major part of pydantic version
+          CURRENT_VERSION=$(poetry run python -c "import pydantic; print(pydantic.__version__)" | cut -d. -f1)
+
+          # Check that the major part of pydantic version is as expected, if not
+          # raise an error
+          if [[ "$CURRENT_VERSION" != "$TEST_WITH_VERSION" ]]; then
+            echo "Error: expected pydantic version ${CURRENT_VERSION} to have been installed, but found: ${TEST_WITH_VERSION}"
+            exit 1
+          fi
+          echo "Found pydantic version ${CURRENT_VERSION}, as expected"
+      - name: Run pydantic compatibility tests
+        shell: bash
+        run: make test

.github/workflows/release.yml

@@ -1,9 +1,7 @@
 name: release
 
 on:
-  pull_request:
-    types:
-      - closed
+  push:
     branches:
       - main
     paths:

kor/_pydantic.py

@@ -0,0 +1,14 @@
+"""Internal code used to support Pydantic v1 and v2."""
+
+
+def _get_pydantic_major_version() -> int:
+    """Get the major version of Pydantic."""
+    try:
+        import pydantic
+
+        return int(pydantic.__version__.split(".")[0])
+    except ImportError:
+        return 0
+
+
+PYDANTIC_MAJOR_VERSION = _get_pydantic_major_version()

kor/adapters.py

@@ -15,6 +15,7 @@
 
 from pydantic import BaseModel
 
+from ._pydantic import PYDANTIC_MAJOR_VERSION
 from .nodes import Bool, ExtractionSchemaNode, Number, Object, Option, Selection, Text
 from .validators import PydanticValidator, Validator
 
@@ -25,6 +26,53 @@
 PRIMITIVE_TYPES = {str, float, int, type(None)}
 
 
+def _is_many(annotation: Any) -> bool:
+    """Determine if the given annotation should map to field many.
+
+    Map to field many if the annotation is a list or a Union where at least one
+    of the arguments is a list type.
+
+    Args:
+        annotation: The annotation to check.
+
+    Returns:
+        bool
+    """
+    origin = get_origin(annotation)
+    if origin is Union:
+        for arg in get_args(annotation):
+            arg_origin = get_origin(arg)
+            if isinstance(arg_origin, type) and issubclass(arg_origin, List):
+                return True
+    if isinstance(origin, type) and issubclass(origin, List):
+        return True
+    return False
+
+
+def _unpack_if_optional_equivalent(annotation: Any) -> Tuple[bool, Any]:
+    """Determine if type is equivalent to an Optional and if so return the inner type.
+
+    Args:
+        annotation: The annotation to check.
+
+    Returns:
+        Tuple[bool, Any]; e.g., Optional[str] -> (True, str)
+    """
+    origin = get_origin(annotation)
+    if origin is Union:
+        args = get_args(annotation)
+        if len(args) == 2 and type(None) in args:
+            if args[0] is type(None):
+                return True, args[1]
+            else:
+                return True, args[0]
+
+    if origin is Optional:
+        return True, get_args(annotation)[0]
+
+    return False, None
+
+
 def _translate_pydantic_to_kor(
     model_class: Type[BaseModel],
     *,
@@ -47,21 +95,35 @@ def _translate_pydantic_to_kor(
     """
 
     attributes: List[Union[ExtractionSchemaNode, Selection, "Object"]] = []
-    for field_name, field in model_class.__fields__.items():
-        field_info = field.field_info
-        extra = field_info.extra
-        if "examples" in extra:
-            field_examples = extra["examples"]
+
+    if PYDANTIC_MAJOR_VERSION == 1:
+        fields_info = model_class.__fields__.items()  # type: ignore[attr-defined]
+    else:
+        fields_info = model_class.model_fields.items()  # type: ignore[attr-defined]
+
+    for field_name, field in fields_info:
+        if PYDANTIC_MAJOR_VERSION == 1:
+            field_info = field.field_info
+            extra = field_info.extra
+            field_examples = extra.get("examples", tuple())
+            field_description = field_info.description or ""
+            type_ = field.type_
         else:
-            field_examples = tuple()
+            type_ = field.annotation
+            field_examples = field.examples or tuple()
+            field_description = field.description or ""
 
-        field_description = field_info.description or ""
+        field_many = _is_many(type_)
+        get_origin(type_)
 
-        type_ = field.type_
-        field_many = get_origin(field.outer_type_) is list
         attribute: Union[ExtractionSchemaNode, Selection, "Object"]
-        # Precedence matters here since bool is a subclass of int
-        if get_origin(type_) is Union:
+
+        is_optional_equivalent, unpacked_optional = _unpack_if_optional_equivalent(
+            type_
+        )
+
+        if get_origin(type_) is Union and not is_optional_equivalent:
+            # Verify that all arguments are primitive types
             args = get_args(type_)
 
             if not all(arg in PRIMITIVE_TYPES for arg in args):
@@ -77,30 +139,42 @@ def _translate_pydantic_to_kor(
                 many=field_many,
             )
         else:
-            if issubclass(type_, BaseModel):
+            # If the type is an Optional or Union equivalent, use the inner type
+            type_to_use = unpacked_optional if is_optional_equivalent else type_
+
+            # If the type is a parameterized generic, we want to extract
+            # the innter type; e.g., List[str] -> str
+            if not isinstance(type_to_use, type):  # i.e., parameterized generic
+                origin_ = get_origin(type_to_use)
+                if not isinstance(origin_, type) or not issubclass(origin_, List):
+                    raise NotImplementedError(f"Unsupported type: {type_to_use}")
+                type_to_use = get_args(type_to_use)[0]  # extract the argument
+
+            if issubclass(type_to_use, BaseModel):
                 attribute = _translate_pydantic_to_kor(
-                    type_,
+                    type_to_use,
                     description=field_description,
                     examples=field_examples,
                     many=field_many,
                     name=field_name,
                 )
-            elif issubclass(type_, bool):
+            # Precedence matters here since bool is a subclass of int
+            elif issubclass(type_to_use, bool):
                 attribute = Bool(
                     id=field_name,
                     examples=field_examples,
                     description=field_description,
                     many=field_many,
                 )
-            elif issubclass(type_, (int, float)):
+            elif issubclass(type_to_use, (int, float)):
                 attribute = Number(
                     id=field_name,
                     examples=field_examples,
                     description=field_description,
                     many=field_many,
                 )
-            elif issubclass(type_, enum.Enum):
-                enum_choices = list(type_)
+            elif issubclass(type_to_use, enum.Enum):
+                enum_choices = list(type_to_use)
                 attribute = Selection(
                     id=field_name,
                     description=field_description,