"""This module contains the schema classes for the medrecord module."""
from __future__ import annotations
from enum import Enum, auto
from typing import (
TYPE_CHECKING,
Dict,
List,
Literal,
Optional,
Tuple,
TypeAlias,
Union,
overload,
)
from medmodels._medmodels import (
PyAttributeDataType,
PyAttributeType,
PyGroupSchema,
PySchema,
PySchemaType,
)
from medmodels.medrecord.datatype import (
DataType,
DateTime,
Duration,
Float,
Int,
Null,
Option,
)
from medmodels.medrecord.datatype import Union as DataTypeUnion
from medmodels.medrecord.types import (
Attributes,
EdgeIndex,
MedRecordAttribute,
NodeIndex,
)
if TYPE_CHECKING:
from medmodels.medrecord.medrecord import MedRecord
from medmodels.medrecord.types import Group
[docs]
class AttributeType(Enum):
"""Enumeration of attribute types."""
Categorical = auto()
Continuous = auto()
Temporal = auto()
Unstructured = auto()
@staticmethod
def _from_py_attribute_type(py_attribute_type: PyAttributeType) -> AttributeType:
"""Converts a PyAttributeType to an AttributeType.
Args:
py_attribute_type (PyAttributeType): The PyAttributeType to convert.
Returns:
AttributeType: The converted AttributeType.
"""
if py_attribute_type == PyAttributeType.Categorical:
return AttributeType.Categorical
if py_attribute_type == PyAttributeType.Continuous:
return AttributeType.Continuous
if py_attribute_type == PyAttributeType.Temporal:
return AttributeType.Temporal
if py_attribute_type == PyAttributeType.Unstructured:
return AttributeType.Unstructured
msg = "Should never be reached"
raise NotImplementedError(msg)
[docs]
@staticmethod
def infer(data_type: DataType) -> AttributeType:
"""Infers the attribute type from the data type.
Args:
data_type (DataType): The data type to infer the attribute type from.
Returns:
AttributeType: The inferred attribute type.
"""
return AttributeType._from_py_attribute_type(
PyAttributeType.infer(data_type._inner())
)
def _into_py_attribute_type(self) -> PyAttributeType:
"""Converts an AttributeType to a PyAttributeType.
Returns:
PyAttributeType: The converted PyAttributeType.
"""
if self == AttributeType.Categorical:
return PyAttributeType.Categorical
if self == AttributeType.Continuous:
return PyAttributeType.Continuous
if self == AttributeType.Temporal:
return PyAttributeType.Temporal
if self == AttributeType.Unstructured:
return PyAttributeType.Unstructured
msg = "Should never be reached"
raise NotImplementedError(msg)
[docs]
def __repr__(self) -> str:
"""Returns a string representation of the AttributeType instance.
Returns:
str: String representation of the attribute type.
"""
return f"AttributeType.{self.name}"
[docs]
def __str__(self) -> str:
"""Returns a string representation of the AttributeType instance.
Returns:
str: String representation of the attribute type.
"""
return self.name
[docs]
def __hash__(self) -> int:
"""Returns the hash of the AttributeType instance.
Returns:
int: The hash of the AttributeType instance.
"""
return hash(self.name)
[docs]
def __eq__(self, value: object) -> bool:
"""Compares the AttributeType instance to another object for equality.
Args:
value (object): The object to compare against.
Returns:
bool: True if the objects are equal, False otherwise.
"""
if isinstance(value, PyAttributeType):
return self._into_py_attribute_type() == value
if isinstance(value, AttributeType):
return str(self) == str(value)
return False
CategoricalType: TypeAlias = DataType
CategoricalPair: TypeAlias = Tuple[CategoricalType, Literal[AttributeType.Categorical]]
ContinuousType: TypeAlias = Union[
Int,
Float,
Null,
Option["ContinuousType"],
DataTypeUnion["ContinuousType", "ContinuousType"],
]
ContinuousPair: TypeAlias = Tuple[ContinuousType, Literal[AttributeType.Continuous]]
TemporalType = Union[
DateTime,
Duration,
Null,
Option["TemporalType"],
DataTypeUnion["TemporalType", "TemporalType"],
]
TemporalPair: TypeAlias = Tuple[TemporalType, Literal[AttributeType.Temporal]]
UnstructuredType: TypeAlias = DataType
UnstructuredPair: TypeAlias = Tuple[
UnstructuredType, Literal[AttributeType.Unstructured]
]
AttributeDataType: TypeAlias = Union[
CategoricalPair, ContinuousPair, TemporalPair, UnstructuredPair
]
AttributesSchema: TypeAlias = Dict[MedRecordAttribute, AttributeDataType]
[docs]
class GroupSchema:
"""A schema for a group of nodes and edges."""
_group_schema: PyGroupSchema
def __init__(
self,
*,
nodes: Optional[
Dict[
MedRecordAttribute,
Union[DataType, AttributeDataType],
],
] = None,
edges: Optional[
Dict[
MedRecordAttribute,
Union[DataType, AttributeDataType],
],
] = None,
) -> None:
"""Initializes a new instance of GroupSchema.
Args:
nodes (Dict[MedRecordAttribute, Union[DataType, AttributeDataType]]):
A dictionary mapping node attributes to their data
types and optional attribute types. Defaults to an empty dictionary.
When no attribute type is provided, it is inferred from the data type.
edges (Dict[MedRecordAttribute, Union[DataType, AttributeDataType]]):
A dictionary mapping edge attributes to their data types and
optional attribute types. Defaults to an empty dictionary.
When no attribute type is provided, it is inferred from the data type.
"""
if edges is None:
edges = {}
if nodes is None:
nodes = {}
def _convert_input(
input: Union[DataType, AttributeDataType],
) -> PyAttributeDataType:
if isinstance(input, tuple):
return PyAttributeDataType(
input[0]._inner(), input[1]._into_py_attribute_type()
)
return PyAttributeDataType(
input._inner(), PyAttributeType.infer(input._inner())
)
self._group_schema = PyGroupSchema(
nodes={x: _convert_input(nodes[x]) for x in nodes},
edges={x: _convert_input(edges[x]) for x in edges},
)
@classmethod
def _from_py_group_schema(cls, group_schema: PyGroupSchema) -> GroupSchema:
"""Creates a GroupSchema instance from an existing PyGroupSchema.
Args:
group_schema (PyGroupSchema): The PyGroupSchema instance to convert.
Returns:
GroupSchema: A new GroupSchema instance.
"""
new_group_schema = cls()
new_group_schema._group_schema = group_schema
return new_group_schema
@property
def nodes(self) -> AttributesSchema:
"""Returns the node attributes in the GroupSchema instance.
Returns:
AttributesSchema: An AttributesSchema object containing the node attributes
and their data types.
"""
def _convert_node(
input: PyAttributeDataType,
) -> AttributeDataType:
# SAFETY: The typing is guaranteed to be correct
return (
DataType._from_py_data_type(input.data_type),
AttributeType._from_py_attribute_type(input.attribute_type),
) # pyright: ignore[reportReturnType]
return {
x: _convert_node(self._group_schema.nodes[x])
for x in self._group_schema.nodes
}
@property
def edges(self) -> AttributesSchema:
"""Returns the edge attributes in the GroupSchema instance.
Returns:
AttributesSchema: An AttributesSchema object containing the edge attributes
and their data types.
"""
def _convert_edge(
input: PyAttributeDataType,
) -> AttributeDataType:
# SAFETY: The typing is guaranteed to be correct
return (
DataType._from_py_data_type(input.data_type),
AttributeType._from_py_attribute_type(input.attribute_type),
) # pyright: ignore[reportReturnType]
return {
x: _convert_edge(self._group_schema.edges[x])
for x in self._group_schema.edges
}
[docs]
def validate_node(self, index: NodeIndex, attributes: Attributes) -> None:
"""Validates the attributes of a node against the schema.
Args:
index (NodeIndex): The index of the node.
attributes (Attributes): The attributes of the node.
"""
self._group_schema.validate_node(index, attributes)
[docs]
def validate_edge(self, index: EdgeIndex, attributes: Attributes) -> None:
"""Validates the attributes of an edge against the schema.
Args:
index (EdgeIndex): The index of the edge.
attributes (Attributes): The attributes of the edge.
"""
self._group_schema.validate_edge(index, attributes)
[docs]
class SchemaType(Enum):
"""Enumeration of schema types."""
Provided = auto()
Inferred = auto()
@staticmethod
def _from_py_schema_type(py_schema_type: PySchemaType) -> SchemaType:
"""Converts a PySchemaType to a SchemaType.
Args:
py_schema_type (PySchemaType): The PySchemaType to convert.
Returns:
SchemaType: The converted SchemaType.
"""
if py_schema_type == PySchemaType.Provided:
return SchemaType.Provided
if py_schema_type == PySchemaType.Inferred:
return SchemaType.Inferred
msg = "Should never be reached"
raise NotImplementedError(msg)
def _into_py_schema_type(self) -> PySchemaType:
"""Converts a SchemaType to a PySchemaType.
Returns:
PySchemaType: The converted PySchemaType.
"""
if self == SchemaType.Provided:
return PySchemaType.Provided
if self == SchemaType.Inferred:
return PySchemaType.Inferred
msg = "Should never be reached"
raise NotImplementedError(msg)
[docs]
class Schema:
"""A schema for a collection of groups."""
_schema: PySchema
def __init__(
self,
*,
groups: Optional[Dict[Group, GroupSchema]] = None,
ungrouped: Optional[GroupSchema] = None,
schema_type: Optional[SchemaType] = None,
) -> None:
"""Initializes a new instance of Schema.
Args:
groups (Dict[Group, GroupSchema], optional): A dictionary of group names
to their schemas. Defaults to None.
ungrouped (Optional[GroupSchema], optional): The group schema for all nodes
not in a group. If not provided, an empty group schema is used.
Defaults to None.
schema_type (Optional[SchemaType], optional): The type of the schema.
If not provided, the schema is of type provided. Defaults to None.
"""
if not ungrouped:
ungrouped = GroupSchema()
if groups is None:
groups = {}
if schema_type:
self._schema = PySchema(
groups={x: groups[x]._group_schema for x in groups},
ungrouped=ungrouped._group_schema,
schema_type=schema_type._into_py_schema_type(),
)
else:
self._schema = PySchema(
groups={x: groups[x]._group_schema for x in groups},
ungrouped=ungrouped._group_schema,
)
[docs]
@classmethod
def infer(cls, medrecord: MedRecord) -> Schema:
"""Infers a schema from a MedRecord instance.
Args:
medrecord (MedRecord): The MedRecord instance to infer the schema from.
Returns:
Schema: The inferred schema.
"""
new_schema = cls()
new_schema._schema = PySchema.infer(medrecord._medrecord)
return new_schema
@classmethod
def _from_py_schema(cls, schema: PySchema) -> Schema:
"""Creates a Schema instance from an existing PySchema.
Args:
schema (PySchema): The PySchema instance to convert.
Returns:
Schema: A new Schema instance.
"""
new_schema = cls()
new_schema._schema = schema
return new_schema
@property
def groups(self) -> List[Group]:
"""Lists all the groups in the Schema instance.
Returns:
List[Group]: A list of groups.
"""
return self._schema.groups
[docs]
def group(self, group: Group) -> GroupSchema:
"""Retrieves the schema for a specific group.
Args:
group (Group): The name of the group.
Returns:
GroupSchema: The schema for the specified group.
Raises:
ValueError: If the group does not exist in the schema.
""" # noqa: DOC502
return GroupSchema._from_py_group_schema(self._schema.group(group))
@property
def ungrouped(self) -> GroupSchema:
"""Retrieves the group schema for all ungrouped nodes and edges.
Returns:
GroupSchema: The ungrouped group schema.
"""
return GroupSchema._from_py_group_schema(self._schema.ungrouped)
@property
def schema_type(self) -> SchemaType:
"""Retrieves the schema type.
Returns:
SchemaType: The schema type.
"""
return SchemaType._from_py_schema_type(self._schema.schema_type)
[docs]
def validate_node(
self, index: NodeIndex, attributes: Attributes, group: Optional[Group] = None
) -> None:
"""Validates the attributes of a node against the schema.
Args:
index (NodeIndex): The index of the node.
attributes (Attributes): The attributes of the node.
group (Optional[Group], optional): The group to validate the node against.
If not provided, the ungrouped schema is used. Defaults to None.
"""
self._schema.validate_node(index, attributes, group)
[docs]
def validate_edge(
self, index: EdgeIndex, attributes: Attributes, group: Optional[Group] = None
) -> None:
"""Validates the attributes of an edge against the schema.
Args:
index (EdgeIndex): The index of the edge.
attributes (Attributes): The attributes of the edge.
group (Optional[Group], optional): The group to validate the edge against.
If not provided, the ungrouped schema is used. Defaults to None.
"""
self._schema.validate_edge(index, attributes, group)
@overload
def set_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[
Literal[AttributeType.Categorical, AttributeType.Unstructured]
] = None,
group: Optional[Group] = None,
) -> None: ...
@overload
def set_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: ContinuousType,
attribute_type: Literal[AttributeType.Continuous],
group: Optional[Group] = None,
) -> None: ...
@overload
def set_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: TemporalType,
attribute_type: Literal[AttributeType.Temporal],
group: Optional[Group] = None,
) -> None: ...
[docs]
def set_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[AttributeType] = None,
group: Optional[Group] = None,
) -> None:
"""Sets the data type and attribute type of a node attribute.
If a data type for the attribute already exists, it is overwritten.
Args:
attribute (MedRecordAttribute): The name of the attribute.
data_type (DataType): The data type of the attribute.
attribute_type (Optional[AttributeType], optional): The attribute type of
the attribute. If not provided, the attribute type is inferred
from the data type. Defaults to None.
group (Optional[Group], optional): The group to set the attribute for.
If no schema for the group exists, a new schema is created.
If not provided, the ungrouped schema is used. Defaults to None.
"""
if not attribute_type:
attribute_type = AttributeType.infer(data_type)
self._schema.set_node_attribute(
attribute,
data_type._inner(),
attribute_type._into_py_attribute_type(),
group,
)
@overload
def set_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[
Literal[AttributeType.Categorical, AttributeType.Unstructured]
] = None,
group: Optional[Group] = None,
) -> None: ...
@overload
def set_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: ContinuousType,
attribute_type: Literal[AttributeType.Continuous],
group: Optional[Group] = None,
) -> None: ...
@overload
def set_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: TemporalType,
attribute_type: Literal[AttributeType.Temporal],
group: Optional[Group] = None,
) -> None: ...
[docs]
def set_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[AttributeType] = None,
group: Optional[Group] = None,
) -> None:
"""Sets the data type and attribute type of an edge attribute.
If a data type for the attribute already exists, it is overwritten.
Args:
attribute (MedRecordAttribute): The name of the attribute.
data_type (DataType): The data type of the attribute.
attribute_type (Optional[AttributeType], optional): The attribute type of
the attribute. If not provided, the attribute type is inferred
from the data type. Defaults to None.
group (Optional[Group], optional): The group to set the attribute for.
If no schema for this group exists, a new schema is created.
If not provided, the ungrouped schema is used. Defaults to None.
"""
if not attribute_type:
attribute_type = AttributeType.infer(data_type)
self._schema.set_edge_attribute(
attribute,
data_type._inner(),
attribute_type._into_py_attribute_type(),
group,
)
@overload
def update_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[
Literal[AttributeType.Categorical, AttributeType.Unstructured]
] = None,
group: Optional[Group] = None,
) -> None: ...
@overload
def update_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: ContinuousType,
attribute_type: Literal[AttributeType.Continuous],
group: Optional[Group] = None,
) -> None: ...
@overload
def update_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: TemporalType,
attribute_type: Literal[AttributeType.Temporal],
group: Optional[Group] = None,
) -> None: ...
[docs]
def update_node_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[AttributeType] = None,
group: Optional[Group] = None,
) -> None:
"""Updates the data type and attribute type of a node attribute.
If a data type for the attribute already exists, it is merged
with the new data type.
Args:
attribute (MedRecordAttribute): The name of the attribute.
data_type (DataType): The data type of the attribute.
attribute_type (Optional[AttributeType], optional): The attribute type of
the attribute. If not provided, the attribute type is inferred
from the data type. Defaults to None.
group (Optional[Group], optional): The group to update the attribute for.
If no schema for this group exists, a new schema is created.
If not provided, the ungrouped schema is used. Defaults to None.
"""
if not attribute_type:
attribute_type = AttributeType.infer(data_type)
self._schema.update_node_attribute(
attribute,
data_type._inner(),
attribute_type._into_py_attribute_type(),
group,
)
@overload
def update_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[
Literal[AttributeType.Categorical, AttributeType.Unstructured]
] = None,
group: Optional[Group] = None,
) -> None: ...
@overload
def update_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: ContinuousType,
attribute_type: Literal[AttributeType.Continuous],
group: Optional[Group] = None,
) -> None: ...
@overload
def update_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: TemporalType,
attribute_type: Literal[AttributeType.Temporal],
group: Optional[Group] = None,
) -> None: ...
[docs]
def update_edge_attribute(
self,
attribute: MedRecordAttribute,
data_type: DataType,
attribute_type: Optional[AttributeType] = None,
group: Optional[Group] = None,
) -> None:
"""Updates the data type and attribute type of an edge attribute.
If a data type for the attribute already exists, it is merged
with the new data type.
Args:
attribute (MedRecordAttribute): The name of the attribute.
data_type (DataType): The data type of the attribute.
attribute_type (Optional[AttributeType], optional): The attribute type of
the attribute. If not provided, the attribute type is inferred
from the data type. Defaults to None.
group (Optional[Group], optional): The group to update the attribute for.
If no schema for this group exists, a new schema is created.
If not provided, the ungrouped schema is used. Defaults to None.
"""
if not attribute_type:
attribute_type = AttributeType.infer(data_type)
self._schema.update_edge_attribute(
attribute,
data_type._inner(),
attribute_type._into_py_attribute_type(),
group,
)
[docs]
def remove_node_attribute(
self, attribute: MedRecordAttribute, group: Optional[Group] = None
) -> None:
"""Removes a node attribute from the schema.
Args:
attribute (MedRecordAttribute): The name of the attribute to remove.
group (Optional[Group], optional): The group to remove the attribute from.
If not provided, the ungrouped schema is used. Defaults to None.
"""
self._schema.remove_node_attribute(attribute, group)
[docs]
def remove_edge_attribute(
self, attribute: MedRecordAttribute, group: Optional[Group] = None
) -> None:
"""Removes an edge attribute from the schema.
Args:
attribute (MedRecordAttribute): The name of the attribute to remove.
group (Optional[Group], optional): The group to remove the attribute from.
If not provided, the ungrouped schema is used. Defaults to None.
"""
self._schema.remove_edge_attribute(attribute, group)
[docs]
def add_group(self, group: Group, group_schema: GroupSchema) -> None:
"""Adds a new group to the schema.
Args:
group (Group): The name of the group.
group_schema (GroupSchema): The schema for the group.
"""
self._schema.add_group(group, group_schema._group_schema)
[docs]
def remove_group(self, group: Group) -> None:
"""Removes a group from the schema.
Args:
group (Group): The name of the group to remove.
"""
self._schema.remove_group(group)
[docs]
def freeze(self) -> None:
"""Freezes the schema. No changes are automatically inferred."""
self._schema.freeze()
[docs]
def unfreeze(self) -> None:
"""Unfreezes the schema. Changes are automatically inferred."""
self._schema.unfreeze()
