improve the metadata to satisfy the need of HMA

sdgx/data_models/combiner.py

@@ -33,8 +33,9 @@ class MetadataCombiner(BaseModel):
 
     relationships: List[Relationship] = []
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, relationships=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.relationships = relationships or []
 
     def check(self):
         """Do necessary checks:

sdgx/data_models/metadata.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
 import json
+import warnings
 from collections import defaultdict
 from collections.abc import Iterable
+from copy import deepcopy
 from itertools import chain
 from pathlib import Path
-from typing import Any, Dict, Set
+from typing import Any, Dict, List, Set
 
 import pandas as pd
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from sdgx.data_loader import DataLoader
 from sdgx.data_models.inspectors.base import RelationshipInspector
@@ -34,6 +36,27 @@ class Metadata(BaseModel):
         column_list(list[str]): list of the comlumn name in the table, other columns lists are used to store column information.
     """
 
+    _SDTYPE_KWARGS = {
+        "numerical": frozenset(["computer_representation"]),
+        "datetime": frozenset(["datetime_format"]),
+        "categorical": frozenset(["order", "order_by"]),
+        "boolean": frozenset([]),
+        "id": frozenset(["regex_format"]),
+        "unknown": frozenset(["pii"]),
+    }
+
+    _KEYS = frozenset(
+        [
+            "columns",
+            "primary_key",
+            "alternate_keys",
+            "sequence_key",
+            "sequence_index",
+            "column_relationships",
+            "METADATA_SPEC_VERSION",
+        ]
+    )
+
     primary_keys: Set[str] = set()
     """
     primary_keys is used to store single primary key or composite primary key
@@ -70,6 +93,157 @@ class Metadata(BaseModel):
     """
     For extend information, use ``get`` and ``set``
     """
+    columns: Dict = defaultdict(str)
+    primary_key: str = "default"
+    alternate_keys: List[str] = Field(default_factory=list, optional=True)
+
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+        self.columns = {}
+        self.primary_key = "default"
+        self.alternate_keys = []
+
+    def add_column(self, column_name, **kwargs):
+        """Add a column to the ``SingleTableMetadata``.
+
+        Args:
+            column_name (str):
+                The column name to be added.
+            kwargs (type):
+                Any additional key word arguments for the column, where ``sdtype`` is required.
+
+        Raises:
+            - ``InvalidMetadataError`` if the column already exists.
+            - ``InvalidMetadataError`` if the ``kwargs`` do not contain ``sdtype``.
+            - ``InvalidMetadataError`` if the column has unexpected values or ``kwargs`` for the
+              given ``sdtype``.
+            - ``InvalidMetadataError`` if the ``pii`` value is not ``True`` or ``False`` when
+               present.
+        """
+        if column_name in self.columns:
+            raise Exception(
+                f"Column name '{column_name}' already exists. Use 'update_column' "
+                "to update an existing column."
+            )
+
+        sdtype = kwargs.get("sdtype")
+        if sdtype is None:
+            raise Exception(f"Please provide a 'sdtype' for column '{column_name}'.")
+
+        column_kwargs = deepcopy(kwargs)
+        if sdtype not in self._SDTYPE_KWARGS:
+            pii = column_kwargs.get("pii", True)
+            column_kwargs["pii"] = pii
+
+        self.columns[column_name] = column_kwargs
+
+    def update_column(self, column_name, **kwargs):
+        """Update an existing column in the ``SingleTableMetadata``.
+
+        Args:
+            column_name (str):
+                The column name to be updated.
+            **kwargs (type):
+                Any key word arguments that describe metadata for the column.
+
+        Raises:
+            - ``InvalidMetadataError`` if the column doesn't already exist in the
+              ``SingleTableMetadata``.
+            - ``InvalidMetadataError`` if the column has unexpected values or ``kwargs`` for the
+              current
+              ``sdtype``.
+            - ``InvalidMetadataError`` if the ``pii`` value is not ``True`` or ``False`` when
+               present.
+        """
+        _kwargs = deepcopy(kwargs)
+        if "sdtype" in kwargs:
+            sdtype = kwargs.pop("sdtype")
+        else:
+            sdtype = self.columns[column_name]["sdtype"]
+            _kwargs["sdtype"] = sdtype
+
+        self.columns[column_name] = _kwargs
+
+    def set_primary_key(self, column_name):
+        """Set the metadata primary key.
+
+        Args:
+            column_name (str):
+                Name of the primary key column(s).
+        """
+        if column_name in self.alternate_keys:
+            warnings.warn(
+                f"'{column_name}' is currently set as an alternate key and will be removed from "
+                "that list."
+            )
+            self.alternate_keys.remove(column_name)
+
+        if self.primary_key is not None:
+            warnings.warn(
+                f"There is an existing primary key '{self.primary_key}'."
+                " This key will be removed."
+            )
+
+        self.primary_key = column_name
+
+    def remove_primary_key(self):
+        """Remove the metadata primary key."""
+        if self.primary_key is None:
+            warnings.warn("No primary key exists to remove.")
+
+        self.primary_key = None
+
+    def add_column_relationship(self, relationship_type, column_names):
+        """Add a column relationship to the metadata.
+
+        Args:
+            relationship_type (str):
+                Type of column relationship.
+            column_names (list[str]):
+                List of column names in the relationship.
+        """
+        relationship = {"type": relationship_type, "column_names": column_names}
+        to_check = [relationship] + self.column_relationships
+        self.column_relationships.append(relationship)
+
+    def _get_primary_and_alternate_keys(self):
+        """Get set of primary and alternate keys.
+
+        Returns:
+            set:
+                Set of keys.
+        """
+        keys = set(self.alternate_keys)
+        if self.primary_key:
+            keys.update({self.primary_key})
+
+        return keys
+
+    @staticmethod
+    def _get_invalid_column_values(column, validation_function):
+        valid = column.apply(validation_function).astype(bool)
+
+        return set(column[~valid])
+
+    @classmethod
+    def load_from_dict(cls, metadata_dict):
+        """Create a ``SingleTableMetadata`` instance from a python ``dict``.
+
+        Args:
+            metadata_dict (dict):
+                Python dictionary representing a ``SingleTableMetadata`` object.
+
+        Returns:
+            Instance of ``SingleTableMetadata``.
+        """
+        instance = cls()
+        for key in instance._KEYS:
+            value = deepcopy(metadata_dict.get(key))
+            if value:
+                setattr(instance, f"{key}", value)
+
+        return instance
+
 
     @property
     def tag_fields(self) -> Iterable[str]: