fix(deps): Make library compatible with Pydantic V2

Author: mikkelkp

poetry.lock

@@ -23,7 +23,7 @@ queenbee = 'queenbee.cli:main'
 
 [tool.poetry.dependencies]
 python = "^3.10"
-pydantic = "<2.0"
+pydantic = "^2.0"
 pyyaml = ">=6.0"
 jsonschema = ">=4.17.3"
 click = "^8.1.7"
@@ -33,7 +33,7 @@ click-plugins = "^1.1.1"
 cli = ["click", "click_plugins"]
 
 [tool.poetry.group.dev.dependencies]
-pydantic-openapi-helper = "^0.2.11"
+pydantic-openapi-helper = "^1.0.2"
 coverage = ">=6"
 coveralls = ">=3"
 pytest = ">=6.2.4"

pyproject.toml

@@ -1,11 +1,11 @@
 """Queenbee utility functions."""
 import hashlib
 import json
-from typing import List, Dict, Any
+from typing import List, Union, Dict, Any
 
 import yaml
 from pydantic import BaseModel as PydanticBaseModel
-from pydantic import validator, Field, constr
+from pydantic import Field, field_validator, ValidationError
 
 from .parser import parse_file
 from .variable import get_ref_variable
@@ -30,37 +30,42 @@ class BaseModelNoType(PydanticBaseModel):
     extensions.
     """
 
-    def yaml(self, exclude_unset=False, **kwargs):
+    def yaml(self, exclude_defaults: bool = False, exclude_none: bool = False, **kwargs):
         """Get a YAML string from the model
 
         Keyword Arguments:
-            exclude_unset {bool} -- Boolean toggle to add or remove any unset/None values
-            (default: {False})
+            exclude_defaults {bool} -- Whether to exclude fields that are set to their
+            default values. (default: {False})
 
         Returns:
             str -- A yaml string representing the model
         """
+        data = self.model_dump(
+            mode='json', by_alias=True, exclude_defaults=exclude_defaults, exclude_none=exclude_none, **kwargs
+        )
         return yaml.dump(
-            json.loads(self.json(by_alias=True,
-                                 exclude_unset=exclude_unset, **kwargs)),
+            data,
             default_flow_style=False
         )
 
-    def to_dict(self, exclude_unset=False, by_alias=True, **kwargs):
+    def to_dict(self, exclude_defaults: bool = False, by_alias: bool = True, exclude_none: bool = False, **kwargs):
         """Get a dictionary from the model
 
         Keyword Arguments:
-            exclude_unset {bool} -- Boolean toggle to add or remove any unset/None values
+            exclude_defaults {bool} -- Whether to exclude fields that are set to their
+                default values.
                 (default: {False})
             by_alias {bool} -- Boolean toggle to use attribute alias or attribute names
                 as key (default: {True})
 
         Returns:
             dict -- A python dictionary representing the model
         """
-        return json.loads(self.json(by_alias=by_alias, exclude_unset=exclude_unset, **kwargs))
+        return self.model_dump(
+            mode='json', by_alias=by_alias, exclude_defaults=exclude_defaults, exclude_none=exclude_none, **kwargs
+        )
 
-    def to_json(self, filepath, indent=None, **kwargs):
+    def to_json(self, filepath: str, indent: int = None, **kwargs):
         """Write a JSON file of the model
 
         Arguments:
@@ -70,20 +75,19 @@ def to_json(self, filepath, indent=None, **kwargs):
             indent {int} -- indent amount (default: {None})
         """
         with open(filepath, 'w') as file:
-            file.write(self.json(by_alias=True, exclude_unset=False,
-                                 indent=indent, **kwargs))
+            file.write(self.model_dump_json(by_alias=True, indent=indent, exclude_none=True, **kwargs))
 
-    def to_yaml(self, filepath, exclude_unset=False, **kwargs):
+    def to_yaml(self, filepath: str, exclude_defaults: bool = False, exclude_none: bool = True, **kwargs):
         """Write a YAML file of the model
 
         Arguments:
             filepath {str} -- Path to the file to be written
 
         Keyword Arguments:
-            exclude_unset {bool} -- Boolean toggle to add or remove any unset/None values
-            (default: {False})
+            exclude_defaults {bool} -- Whether to exclude fields that are set to their
+            default values. (default: {False})
         """
-        content = self.yaml(exclude_unset=exclude_unset, **kwargs)
+        content = self.yaml(exclude_defaults=exclude_defaults, exclude_none=exclude_none, **kwargs)
 
         with open(filepath, 'w') as out_file:
             out_file.write(content)
@@ -99,7 +103,10 @@ def from_file(cls, filepath):
             cls -- An instance of the pydantic class
         """
         data = parse_file(filepath)
-        return cls.parse_obj(data)
+        try:
+            return cls.model_validate(data)
+        except ValidationError as e:
+            raise ValueError(e) from e
 
     def __repr__(self):
         return self.yaml()
@@ -112,7 +119,7 @@ def __hash__(self):
             str -- A hash/digest of the model
         """
         return hashlib.sha256(
-            self.json(by_alias=True, exclude_unset=False).encode('utf-8')
+            self.model_dump_json(by_alias=True, exclude_none=True).encode('utf-8')
         ).hexdigest()
 
     def _referenced_values(self, var_names: List[str]) -> Dict[str, List[str]]:
@@ -143,16 +150,17 @@ def _referenced_values(self, var_names: List[str]) -> Dict[str, List[str]]:
 class BaseModel(BaseModelNoType):
     """BaseModel with functionality to return the object as a yaml string."""
 
-    type: constr(regex='^BaseModel$') = 'BaseModel'
+    type: str = Field('BaseModel', pattern='^BaseModel$')
 
     annotations: Dict[str, Any] = Field(
-        None,
+        default_factory=dict,
         description='An optional dictionary to add annotations to inputs. These '
         'annotations will be used by the client side libraries.'
     )
 
-    @validator('type')
-    def ensure_type_match(cls, v):
+    @field_validator('type')
+    @classmethod
+    def ensure_type_match(cls, v: str) -> str:
         name = cls.__name__
         if name != v:
             raise ValueError(
@@ -161,7 +169,3 @@ def ensure_type_match(cls, v):
                 f'Pydantic object. In this case: {name}'
             )
         return v
-
-    @validator('annotations', always=True)
-    def replace_none_value(cls, v):
-        return {} if not v else v

queenbee/base/basemodel.py

@@ -6,45 +6,45 @@
 https://swagger.io/specification/#infoObject
 """
 
-from typing import List
-from pydantic import Field, constr, AnyUrl
+from typing import List, Literal, Union
+from pydantic import Field, AnyUrl
 
 from .basemodel import BaseModel
 
 
 class Maintainer(BaseModel):
     """Maintainer information"""
-    type: constr(regex='^Maintainer$') = 'Maintainer'
+    type: Literal['Maintainer'] = 'Maintainer'
 
     name: str = Field(
         ...,
         description='The name of the author/maintainer person or organization.'
     )
 
-    email: str = Field(
+    email: Union[str, None] = Field(
         None,
         description='The email address of the author/maintainer person or organization.'
     )
 
 
 class License(BaseModel):
     """License information for the Package"""
-    type: constr(regex='^License$') = 'License'
+    type: Literal['License'] = 'License'
 
     name: str = Field(
         ...,
         description='The license name used for the package.'
     )
 
-    url: AnyUrl = Field(
+    url: Union[AnyUrl, None] = Field(
         None,
         description='A URL to the license used for the package.'
     )
 
 
 class MetaData(BaseModel):
     """Package metadata information."""
-    type: constr(regex='^MetaData$') = 'MetaData'
+    type: Literal['MetaData'] = 'MetaData'
 
     name: str = Field(
         ...,
@@ -56,47 +56,47 @@ class MetaData(BaseModel):
         description='The tag of the package'
     )
 
-    app_version: str = Field(
+    app_version: Union[str, None] = Field(
         None,
         description='The version of the application code underlying the manifest'
     )
 
-    keywords: List[str] = Field(
+    keywords: Union[List[str], None] = Field(
         None,
         description='A list of keywords to search the package by'
     )
 
-    maintainers: List[Maintainer] = Field(
+    maintainers: Union[List[Maintainer], None] = Field(
         None,
         description='A list of maintainers for the package'
     )
 
-    home: str = Field(
+    home: Union[str, None] = Field(
         None,
         description='The URL of this package\'s home page'
     )
 
-    sources: List[str] = Field(
+    sources: Union[List[str], None] = Field(
         None,
         description='A list of URLs to source code for this project'
     )
 
-    icon: str = Field(
+    icon: Union[str, None] = Field(
         None,
         description='A URL to an SVG or PNG image to be used as an icon'
     )
 
-    deprecated: bool = Field(
+    deprecated: Union[bool, None] = Field(
         None,
         description='Whether this package is deprecated'
     )
 
-    description: str = Field(
+    description: Union[str, None] = Field(
         None,
         description='A description of what this package does'
     )
 
-    license: License = Field(
+    license: Union[License, None] = Field(
         None,
         description='The license information.'
     )

queenbee/base/metadata.py

@@ -74,7 +74,7 @@ def new(name, path):
     input_dict = parse_file(path)
     input_dict['metadata']['name'] = name
 
-    plugin = Plugin.parse_obj(input_dict)
+    plugin = Plugin.model_validate(input_dict)
 
     # Readme
     readme_string = f"""

queenbee/cli/plugin/new.py

@@ -42,7 +42,7 @@ def lint(path, dependency_update):
             config=ctx.obj.config
         )
         obj = br.to_dict()
-        BakedRecipe.parse_obj(obj)
+        BakedRecipe.model_validate(obj)
     except ValidationError as error:
         raise click.ClickException(error)
     except FileNotFoundError as error:

queenbee/cli/recipe/lint.py

@@ -82,7 +82,7 @@ def new(name, path):
     input_dict = parse_file(path)
     input_dict['metadata']['name'] = name
 
-    recipe = Recipe.parse_obj(input_dict)
+    recipe = Recipe.model_validate(input_dict)
 
     # Readme
     readme_string = f"""

queenbee/cli/recipe/new.py

@@ -1,14 +1,14 @@
-from typing import List, Union, Dict
+from typing import List, Union, Dict, Literal
 from urllib.parse import urlparse
-from pydantic import Field, SecretStr, constr
+from pydantic import Field, SecretStr
 
 from ..base.basemodel import BaseModel
 from .auth import JWTAuth, HeaderAuth
 from .repositories import RepositoryReference
 
 
 class Config(BaseModel):
-    type: constr(regex='^Config$') = 'Config'
+    type: Literal['Config'] = 'Config'
 
     auth: List[Union[JWTAuth, HeaderAuth]] = Field(
         [],
@@ -111,8 +111,3 @@ def remove_repository(self, name: str):
                 existing_index = i
 
         del self.repositories[existing_index]
-
-    class Config:
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value() if v else None,
-        }

queenbee/config/__init__.py

@@ -1,28 +1,21 @@
-import json
-from enum import Enum
-from urllib import request
-from urllib.error import HTTPError
-from typing import Union, Dict
-from pydantic import Field, SecretStr, constr
-
+from typing import Dict, Literal, Union
+from pydantic import Field, SecretStr
 
 from ..base.basemodel import BaseModel
-from ..base.request import make_request, USER_AGENT_STRING
 
 
 class BaseAuth(BaseModel):
-
-    type: constr(regex='^BaseAuth$') = 'BaseAuth'
+    type: Literal['BaseAuth'] = 'BaseAuth'
 
     domain: str = Field(
         ...,
         description='The host domain to authenticate to',
-        example='api.pollination.solutions'
+        json_schema_extra={'example': 'api.pollination.solutions'}
     )
 
-    access_token: SecretStr = Field(
+    access_token: Union[SecretStr, None] = Field(
         None,
-        description='A JWT token retrieved from a previous login'
+        description='An access token for the domain.'
     )
 
     @property
@@ -41,12 +34,11 @@ def refresh_token(self):
 
 
 class HeaderAuth(BaseAuth):
-
-    type: Enum('JWTAuth', {'type': 'jwt'}) = 'jwt'
+    type: Literal['HeaderAuth'] = 'HeaderAuth'
 
     header_name: str = Field(
         ...,
-        description='The HTTP header to user'
+        description='The HTTP header to use'
     )
 
     @property
@@ -56,11 +48,10 @@ def auth_header(self) -> Dict[str, str]:
         Returns:
             Dict[str, str]: a header with an API token
         """
+        if self.access_token is None:
+            return {}
         return {self.header_name: self.access_token.get_secret_value()}
 
 
 class JWTAuth(BaseAuth):
-
-    type: Enum('JWTAuth', {'type': 'jwt'}) = 'jwt'
-
-
+    type: Literal['JWTAuth'] = 'JWTAuth'