Skip to content

Commit

Permalink
fix: DocVec display (docarray#1522)
Browse files Browse the repository at this point in the history
Signed-off-by: AnneY <evangeline-lun@foxmail.com>
  • Loading branch information
AnneYang720 authored May 10, 2023
1 parent 096f644 commit bfebad6
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 2 deletions.
3 changes: 3 additions & 0 deletions docarray/array/doc_vec/column_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,3 +140,6 @@ def __iter__(self):

def __len__(self):
return len(self.storage.columns)

def keys(self):
return self.storage.columns.keys()
11 changes: 10 additions & 1 deletion docarray/base_doc/doc.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,8 +104,17 @@ def _get_field_type(cls, field: str) -> Type:
return cls.__fields__[field].outer_type_

def __str__(self) -> str:
content: Any = None
if self.is_view():
attr_str = ", ".join(
f"{field}={self.__getattr__(field)}" for field in self.__dict__.keys()
)
content = f"{self.__class__.__name__}({attr_str})"
else:
content = self

with _console.capture() as capture:
_console.print(self)
_console.print(content)

return capture.get().strip()

Expand Down
9 changes: 8 additions & 1 deletion tests/units/document/test_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,14 @@ class MyDoc(BaseDoc):

docs = [MyDoc(tensor=np.zeros((10, 10)), name='hello', id=i) for i in range(4)]

storage = DocVec[MyDoc](docs)._storage
doc_vec = DocVec[MyDoc](docs)
storage = doc_vec._storage

result = str(doc_vec[0])
assert 'MyDoc' in result
assert 'id' in result
assert 'tensor' in result
assert 'name' in result

doc = MyDoc.from_view(ColumnStorageView(0, storage))
assert doc.is_view()
Expand Down

0 comments on commit bfebad6

Please sign in to comment.