Support Pydantic v2 environments via falling back to v1 all of the time (#415)

* Support Pydantic v2

* Bump and CHANGELOG

Author: gatli

CHANGELOG.md

@@ -5,16 +5,11 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [0.16.12](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.12) - 2023-11-27
-
-### Added
-
-- Added `num_processes` parameter to `dataset.items_and_annotation_chip_generator()` to specify parallel processing.
-
 ## [0.16.11](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.11) - 2023-11-22
 
 ### Added
 
+- Added `num_processes` parameter to `dataset.items_and_annotation_chip_generator()` to specify parallel processing.
 - Method to allow for concurrent task fetches for pointcloud data
 
 Example:
@@ -28,6 +23,9 @@ Example:
 }
 ```
 
+### Fixes
+- Support environments using pydantic>=2
+
 ## [0.16.10](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.16.10) - 2023-11-22
 
 Allow creating a dataset by crawling all images in a directory, recursively. Also supports privacy mode datasets.

nucleus/__init__.py

@@ -47,7 +47,12 @@
 import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-import pydantic
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    import pydantic.v1 as pydantic
+except ImportError:
+    import pydantic
+
 import requests
 import tqdm
 

nucleus/data_transfer_object/job_status.py

@@ -4,7 +4,12 @@
 from typing import List, Optional, Union
 
 from dateutil.parser import ParserError, parse
-from pydantic import validator
+
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import validator
+except ImportError:
+    from pydantic import validator
 
 from nucleus.constants import JOB_REQ_LIMIT
 from nucleus.job import CustomerJobTypes

nucleus/pydantic_base.py

@@ -1,4 +1,14 @@
-from pydantic import BaseModel  # pylint: disable=no-name-in-module
+"""
+NOTE:
+We started using pydantic during v1 and are kind of stuck with it now unless we write a compatibility layers.
+As a library we want to support v1 and v2 such that we're not causing downstream problems for our users.
+This means we have to do some import shenanigans to support both v1 and v2.
+"""
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import BaseModel  # pylint: disable=no-name-in-module
+except ImportError:
+    from pydantic import BaseModel
 
 
 class ImmutableModel(BaseModel):

nucleus/test_launch_integration.py

@@ -2,71 +2,101 @@
 from typing import Any, Callable, Dict, List, Optional, Type
 
 from PIL import Image, ImageDraw
-from pydantic import BaseModel, Extra, ValidationError
+
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import BaseModel, Extra, ValidationError
+except ImportError:
+    from pydantic import BaseModel, Extra, ValidationError
 
 # From scaleapi/server/src/lib/select/api/types.ts
 # These classes specify how user models must pass output to Launch + Nucleus.
 
 
-class PointModel(BaseModel, extra=Extra.forbid):
+class PointModel(BaseModel):
     x: float
     y: float
 
+    class Config:
+        extra = Extra.forbid
+
 
-class BoxGeometryModel(BaseModel, extra=Extra.forbid):
+class BoxGeometryModel(BaseModel):
     x: float
     y: float
     width: float
     height: float
 
+    class Config:
+        extra = Extra.forbid
+
 
-class BoxAnnotationModel(BaseModel, extra=Extra.forbid):
+class BoxAnnotationModel(BaseModel):
     geometry: BoxGeometryModel
     type: str
     label: Optional[str] = None
     confidence: Optional[float] = None
     classPdf: Optional[Dict[str, float]] = None
     metadata: Optional[Dict[str, Any]] = None
 
+    class Config:
+        extra = Extra.forbid
 
-class NoneGeometryModel(BaseModel, extra=Extra.forbid):
-    pass
 
+class NoneGeometryModel(BaseModel):
+    class Config:
+        extra = Extra.forbid
 
-class CategoryAnnotationModel(BaseModel, extra=Extra.forbid):
+
+class CategoryAnnotationModel(BaseModel):
     geometry: NoneGeometryModel
     type: str
     label: Optional[str] = None
     confidence: Optional[float] = None
     classPdf: Optional[Dict[str, float]] = None
     metadata: Optional[Dict[str, Any]] = None
 
+    class Config:
+        extra = Extra.forbid
+
 
-class LineGeometryModel(BaseModel, extra=Extra.forbid):
+class LineGeometryModel(BaseModel):
     vertices: List[PointModel]
 
+    class Config:
+        extra = Extra.forbid
 
-class LineAnnotationModel(BaseModel, extra=Extra.forbid):
+
+class LineAnnotationModel(BaseModel):
     geometry: LineGeometryModel
     type: str
     label: Optional[str] = None
     confidence: Optional[float] = None
     classPdf: Optional[Dict[str, float]] = None
     metadata: Optional[Dict[str, Any]] = None
 
+    class Config:
+        extra = Extra.forbid
+
 
-class PolygonGeometryModel(BaseModel, extra=Extra.forbid):
+class PolygonGeometryModel(BaseModel):
     vertices: List[PointModel]
 
+    class Config:
+        extra = Extra.forbid
 
-class PolygonAnnotationModel(BaseModel, extra=Extra.forbid):
+
+class PolygonAnnotationModel(BaseModel):
     geometry: PolygonGeometryModel
     type: str
     label: Optional[str] = None
     confidence: Optional[float] = None
     classPdf: Optional[Dict[str, float]] = None
     metadata: Optional[Dict[str, Any]] = None
 
+    class Config:
+        extra = Extra.forbid
+
 
 def verify_output(
     annotation_list: List[Dict[str, Any]],

nucleus/validate/data_transfer_objects/eval_function.py

@@ -1,6 +1,10 @@
 from typing import Any, Dict, List, Optional
 
-from pydantic import validator
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import validator
+except ImportError:
+    from pydantic import validator
 
 from ...pydantic_base import ImmutableModel
 from ..constants import ThresholdComparison

nucleus/validate/data_transfer_objects/scenario_test.py

@@ -1,6 +1,10 @@
 from typing import List
 
-from pydantic import validator
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import validator
+except ImportError:
+    from pydantic import validator
 
 from nucleus.pydantic_base import ImmutableModel
 

nucleus/validate/data_transfer_objects/scenario_test_evaluations.py

@@ -1,6 +1,10 @@
 from typing import Optional
 
-from pydantic import root_validator, validator
+try:
+    # NOTE: we always use pydantic v1 but have to do these shenanigans to support both v1 and v2
+    from pydantic.v1 import root_validator, validator
+except ImportError:
+    from pydantic import root_validator, validator
 
 from nucleus.pydantic_base import ImmutableModel
 

pyproject.toml

@@ -25,7 +25,7 @@ ignore = ["E501", "E741", "E731", "F401"]  # Easy ignore for getting it running
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.16.10"
+version = "0.16.11"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]