Code_Versions/workflow/io_port.py

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Dict, Any, List
# Codes Go Below:


# Exceptions
class CompatibilityException(Exception):
    def __init__(self, reason: str):
        super().__init__(reason)
        self.reason = reason


# Base DType (Semantic Contract)
class DType(ABC):
    """
    Pure semantic type. Immutable by design.
    """

    def __init__(self, metadata: Optional[Dict[str, Any]] = None):
        self._metadata = metadata or {}

    @property
    def metadata(self):
        return self._metadata

    @abstractmethod
    def can_connect_to(self, other: "DType"):
        pass

    def __repr__(self):
        return f"{self.__class__.__name__}(metadata={self.metadata})"


# Composite Types (Immutable)
@dataclass(frozen=True)
class CompositePart:
    name: str
    dtype: DType


@dataclass(frozen=True)
class CompositeType(DType):
    parts: List[CompositePart]
    metadata: Optional[Dict[str, Any]] = None

    def __post_init__(self):
        object.__setattr__(self, "_metadata", self.metadata or {})

    def can_connect_to(self, other: "DType"):
        if not isinstance(other, CompositeType):
            raise CompatibilityException("Expected CompositeType")

        if len(self.parts) != len(other.parts):
            raise CompatibilityException("CompositeType length mismatch")

        for p1, p2 in zip(self.parts, other.parts):
            if p1.name != p2.name:
                raise CompatibilityException(
                    f"CompositeType part mismatch: {p1.name} vs {p2.name}"
                )
            p1.dtype.can_connect_to(p2.dtype)


# Image Types
class ImageType(DType):
    def __init__(self, modality: Optional[str] = None, metadata=None):
        super().__init__(metadata)
        self.modality = modality

    def can_connect_to(self, other: "DType"):
        if not isinstance(other, ImageType):
            raise CompatibilityException("Target is not ImageType")

        if self.modality and other.modality and self.modality != other.modality:
            raise CompatibilityException(
                f"Image modality mismatch: {self.modality} vs {other.modality}"
            )


class NIFTIImageType(ImageType):
    def __init__(self, metadata=None):
        super().__init__("nifti", metadata)


class DICOMImageType(ImageType):
    def __init__(self, metadata=None):
        super().__init__("dicom", metadata)


# Table Types
class TableType(DType):
    def __init__(self, columns: List[str], metadata=None):
        super().__init__(metadata)
        self.columns = columns

    def can_connect_to(self, other: "DType"):
        if not isinstance(other, TableType):
            raise CompatibilityException("Target is not TableType")

        if self.columns != other.columns:
            raise CompatibilityException(
                f"Table schema mismatch: {self.columns} vs {other.columns}"
            )


class CSVTableType(TableType):
    def __init__(self, columns: List[str], delimiter=",", encoding="utf-8", metadata=None):
        super().__init__(columns, metadata)
        self.delimiter = delimiter
        self.encoding = encoding

    def can_connect_to(self, other: "DType"):
        super().can_connect_to(other)

        if not isinstance(other, CSVTableType):
            raise CompatibilityException("CSV tables must match CSV tables")

        if self.delimiter != other.delimiter:
            raise CompatibilityException(
                f"CSV delimiter mismatch: {self.delimiter} vs {other.delimiter}"
            )


# Ports (Pure Model)
class Port(ABC):
    def __init__(self, name: str, dtype: DType):
        self._name = name
        self._dtype = dtype
        self._parent_task = None

    @property
    def name(self):
        return self._name

    @property
    def dtype(self):
        return self._dtype

    @property
    def parent_task(self):
        return self._parent_task

    def full_name(self):
        if self.parent_task:
            return f"{self.parent_task.name}.{self.name}"
        return self.name


# OutPort
class OutPort(Port):
    def __init__(self, name: str, dtype: DType):
        super().__init__(name, dtype)
        self._connections: List["InPort"] = []

    def connect(self, in_port: "InPort"):
        in_port.dtype.can_connect_to(self.dtype)
        self._connections.append(in_port)
        in_port._connected_output = self

    @property
    def connections(self):
        return list(self._connections)


# InPort
class InPort(Port):
    def __init__(self, name: str, dtype: DType, required=True):
        super().__init__(name, dtype)
        self.required = required
        self._connected_output: Optional[OutPort] = None

    @property
    def connected_output(self):
        return self._connected_output

    def is_connected(self):
        return self._connected_output is not None

    def validate(self):
        if self.required and not self.is_connected():
            raise CompatibilityException(
                f"InPort {self.full_name()} is required but not connected"
            )