Fix/column annotation (#184)

* content type validator and serializer

* add pdb content type

Author: zhu0619

polaris/dataset/_column.py

@@ -18,7 +18,7 @@ class Modality(enum.Enum):
     IMAGE = "image"
 
 
-class KnownContentType(str, enum.Enum):
+class KnownContentType(enum.Enum):
     """Used to specify column's IANA content type in a dataset."""
 
     SMILES = "chemical/x-smiles"
@@ -57,6 +57,13 @@ def _validate_modality(cls, v, values):
             v = Modality[v.upper()]
         return v
 
+    @field_validator("content_type")
+    def _validate_content_type(cls, v, values):
+        """Tries to convert a string to the Enum"""
+        if isinstance(v, str):
+            v = KnownContentType[v.upper()]
+        return v
+
     @field_validator("dtype")
     def _validate_dtype(cls, v):
         """Tries to convert a string to the Enum"""
@@ -69,6 +76,13 @@ def _serialize_modality(self, v: Modality):
         """Return the modality as a string, keeping it serializable"""
         return v.name
 
+    @field_serializer("content_type")
+    def _serialize_content_type(self, v: KnownContentType):
+        """Return the content_type as a string, keeping it serializable"""
+        if v is not None:
+            v = v.name
+        return v
+
     @field_serializer("dtype")
     def _serialize_dtype(self, v: Optional[DTypeLike]):
         """Return the dtype as a string, keeping it serializable"""

polaris/dataset/converters/_pdb.py

@@ -7,7 +7,7 @@
 import zarr
 from fastpdb import struc
 
-from polaris.dataset import ColumnAnnotation, Modality
+from polaris.dataset import ColumnAnnotation, Modality, KnownContentType
 from polaris.dataset._adapters import Adapter
 from polaris.dataset.converters._base import Converter, FactoryProduct
 
@@ -188,7 +188,11 @@ def convert(self, path, factory: "DatasetFactory", append: bool = False) -> Fact
         df[self.pdb_column] = pd.Series(pointers)
 
         # Set the annotations
-        annotations = {self.pdb_column: ColumnAnnotation(is_pointer=True, modality=Modality.PROTEIN_3D)}
+        annotations = {
+            self.pdb_column: ColumnAnnotation(
+                is_pointer=True, modality=Modality.PROTEIN_3D, content_type=KnownContentType.PDB
+            )
+        }
 
         # Return the dataframe and the annotations
         return df, annotations, {self.pdb_column: Adapter.ARRAY_TO_PDB}