Skip to content

Commit 661f8d7

Browse files
authored
feat: implemented _str_ and _repr_ methods for python schema classes (#1095)
* feat: implemented _str_ and _repr_ methods for python schema classes * fix: improve string formatting in BasicValueType and TableType classes * fix: remove UTable from TABLE_TYPES and update TableType class to reflect changes
1 parent 1602635 commit 661f8d7

File tree

1 file changed

+75
-0
lines changed

1 file changed

+75
-0
lines changed

python/cocoindex/typing.py

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -516,6 +516,13 @@ class VectorTypeSchema:
516516
element_type: "BasicValueType"
517517
dimension: int | None
518518

519+
def __str__(self) -> str:
520+
dimension_str = f", {self.dimension}" if self.dimension is not None else ""
521+
return f"Vector[{self.element_type}{dimension_str}]"
522+
523+
def __repr__(self) -> str:
524+
return self.__str__()
525+
519526
@staticmethod
520527
def decode(obj: dict[str, Any]) -> "VectorTypeSchema":
521528
return VectorTypeSchema(
@@ -534,6 +541,13 @@ def encode(self) -> dict[str, Any]:
534541
class UnionTypeSchema:
535542
variants: list["BasicValueType"]
536543

544+
def __str__(self) -> str:
545+
types_str = " | ".join(str(t) for t in self.variants)
546+
return f"Union[{types_str}]"
547+
548+
def __repr__(self) -> str:
549+
return self.__str__()
550+
537551
@staticmethod
538552
def decode(obj: dict[str, Any]) -> "UnionTypeSchema":
539553
return UnionTypeSchema(
@@ -573,6 +587,23 @@ class BasicValueType:
573587
vector: VectorTypeSchema | None = None
574588
union: UnionTypeSchema | None = None
575589

590+
def __str__(self) -> str:
591+
if self.kind == "Vector" and self.vector is not None:
592+
dimension_str = (
593+
f", {self.vector.dimension}"
594+
if self.vector.dimension is not None
595+
else ""
596+
)
597+
return f"Vector[{self.vector.element_type}{dimension_str}]"
598+
elif self.kind == "Union" and self.union is not None:
599+
types_str = " | ".join(str(t) for t in self.union.variants)
600+
return f"Union[{types_str}]"
601+
else:
602+
return self.kind
603+
604+
def __repr__(self) -> str:
605+
return self.__str__()
606+
576607
@staticmethod
577608
def decode(obj: dict[str, Any]) -> "BasicValueType":
578609
kind = obj["kind"]
@@ -603,6 +634,18 @@ class EnrichedValueType:
603634
nullable: bool = False
604635
attrs: dict[str, Any] | None = None
605636

637+
def __str__(self) -> str:
638+
result = str(self.type)
639+
if self.nullable:
640+
result += "?"
641+
if self.attrs:
642+
attrs_str = ", ".join(f"{k}: {v}" for k, v in self.attrs.items())
643+
result += f" [{attrs_str}]"
644+
return result
645+
646+
def __repr__(self) -> str:
647+
return self.__str__()
648+
606649
@staticmethod
607650
def decode(obj: dict[str, Any]) -> "EnrichedValueType":
608651
return EnrichedValueType(
@@ -625,6 +668,12 @@ class FieldSchema:
625668
name: str
626669
value_type: EnrichedValueType
627670

671+
def __str__(self) -> str:
672+
return f"{self.name}: {self.value_type}"
673+
674+
def __repr__(self) -> str:
675+
return self.__str__()
676+
628677
@staticmethod
629678
def decode(obj: dict[str, Any]) -> "FieldSchema":
630679
return FieldSchema(name=obj["name"], value_type=EnrichedValueType.decode(obj))
@@ -640,6 +689,13 @@ class StructSchema:
640689
fields: list[FieldSchema]
641690
description: str | None = None
642691

692+
def __str__(self) -> str:
693+
fields_str = ", ".join(str(field) for field in self.fields)
694+
return f"Struct({fields_str})"
695+
696+
def __repr__(self) -> str:
697+
return self.__str__()
698+
643699
@classmethod
644700
def decode(cls, obj: dict[str, Any]) -> Self:
645701
return cls(
@@ -658,6 +714,13 @@ def encode(self) -> dict[str, Any]:
658714
class StructType(StructSchema):
659715
kind: Literal["Struct"] = "Struct"
660716

717+
def __str__(self) -> str:
718+
# Use the parent's __str__ method for consistency
719+
return super().__str__()
720+
721+
def __repr__(self) -> str:
722+
return self.__str__()
723+
661724
def encode(self) -> dict[str, Any]:
662725
result = super().encode()
663726
result["kind"] = self.kind
@@ -670,6 +733,18 @@ class TableType:
670733
row: StructSchema
671734
num_key_parts: int | None = None # Only for KTable
672735

736+
def __str__(self) -> str:
737+
if self.kind == "KTable":
738+
num_parts = self.num_key_parts if self.num_key_parts is not None else 1
739+
table_kind = f"KTable({num_parts})"
740+
else: # LTable
741+
table_kind = "LTable"
742+
743+
return f"{table_kind}({self.row})"
744+
745+
def __repr__(self) -> str:
746+
return self.__str__()
747+
673748
@staticmethod
674749
def decode(obj: dict[str, Any]) -> "TableType":
675750
row_obj = obj["row"]

0 commit comments

Comments
 (0)