addd pydantic dataclasses

Author: Gabefire

libs/labelbox/src/labelbox/orm/model.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, Dict, List, Type, Union
 

libs/labelbox/src/labelbox/schema/export_task.py

@@ -2,7 +2,7 @@
 import os
 import tempfile
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 from enum import Enum
 from functools import lru_cache
 from typing import (

libs/labelbox/src/labelbox/schema/iam_integration.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 
 from labelbox.utils import snake_case
 from labelbox.orm.db_object import DbObject

libs/labelbox/src/labelbox/schema/invite.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field

libs/labelbox/src/labelbox/schema/ontology.py

@@ -2,7 +2,8 @@
 
 import colorsys
 import json
-from dataclasses import dataclass, field
+from pydantic.dataclasses import dataclass
+from dataclasses import field
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 

libs/labelbox/src/labelbox/schema/role.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 from typing import Dict, Optional, TYPE_CHECKING
 
 from labelbox.orm.model import Field

libs/labelbox/src/labelbox/schema/slice.py

@@ -1,5 +1,7 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
+from pydantic import ConfigDict, GetCoreSchemaHandler
+from pydantic.dataclasses import dataclass
+from pydantic_core import core_schema
+from typing import Any, Optional, Tuple, Union
 import warnings
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field
@@ -47,6 +49,35 @@ def to_hash(self):
                 "global_key": self.global_key.key if self.global_key else None,
             }
 
+        model_config = ConfigDict(arbitrary_types_allowed=True)
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            self, source: type[Any], handler: GetCoreSchemaHandler
+        ) -> core_schema.CoreSchema:
+            print(source)
+            return core_schema.no_info_after_validator_function(
+                self._validate,
+                core_schema.dict_schema(),
+                serialization=core_schema.plain_serializer_function_ser_schema(
+                    lambda x: x.to_hash(),
+                    return_schema=core_schema.dict_schema(),
+                ),
+            )
+
+        @classmethod
+        def _validate(cls, value: Any) -> Any:
+            if isinstance(value, cls):
+                return value
+            if isinstance(value, dict):
+                id_val = value.get("id")
+                if id_val is None:
+                    raise ValueError(
+                        f"Missing required 'id' field in dict {value}"
+                    )
+                return cls(id=id_val, global_key=value.get("global_key"))
+            raise ValueError(f"Cannot convert {value} to {cls.__name__}")
+
 
 class CatalogSlice(Slice):
     """

libs/labelbox/src/labelbox/schema/tool_building/base_step_reasoning_tool.py

@@ -1,6 +1,7 @@
 import warnings
 from abc import ABC
-from dataclasses import dataclass, field
+from pydantic.dataclasses import dataclass
+from dataclasses import field
 from typing import Any, Dict, List, Optional, Set
 
 from labelbox.schema.tool_building.tool_type import ToolType

libs/labelbox/src/labelbox/schema/tool_building/classification.py

@@ -1,5 +1,6 @@
 import warnings
-from dataclasses import dataclass, field
+from pydantic.dataclasses import dataclass
+from dataclasses import field
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 

libs/labelbox/src/labelbox/schema/tool_building/fact_checking_tool.py

@@ -1,4 +1,5 @@
-from dataclasses import dataclass, field
+from pydantic.dataclasses import dataclass
+from dataclasses import field
 from enum import Enum
 
 from labelbox.schema.tool_building.base_step_reasoning_tool import (