Beautiful print of Container base class

Fix #67

easyidp/__init__.py

@@ -15,16 +15,17 @@ class Container(dict):
 
     Caution
     -------
-    This object can not be saved by ``pickle``, it will cause problem when loading
+    This object can not be saved by ``pickle``, it will cause problem when loading [1]_ 
 
     References
     ----------
-    https://stackoverflow.com/questions/4014621/a-python-class-that-acts-like-dict
+    [1] https://stackoverflow.com/questions/4014621/a-python-class-that-acts-like-dict
     """
     def __init__(self):
         super().__init__()
         self.id_item = {}   # {0: item1, 1: item2}
         self.item_label = {}  #{"N1W1": 0, "N1W2": 1}, just index it position
+        self._btf_print = "<Empty easyidp.Container object>"
 
     def __setitem__(self, key, item):
         if isinstance(key, int):
@@ -50,6 +51,8 @@ def __setitem__(self, key, item):
         else:
             raise KeyError(f"Key should be 'int', 'str', not {key}")
 
+        self._update_btf_print()
+
     def __getitem__(self, key):
         if isinstance(key, int):  # index by photo order
             return self.id_item[key]
@@ -61,13 +64,41 @@ def __getitem__(self, key):
             out = self.copy()
             out.id_item = {k:v for k, v in self.id_item.items() if k in idx_list}
             out.item_label = {k:v for k, v in self.item_label.items() if v in idx_list}
+
+            out._update_btf_print()
 
             return out
         else:
             raise KeyError(f"Key should be 'int', 'str', 'slice', not {key}")
 
-    def __repr__(self):
-        return repr(self.id_item)
+    def __repr__(self) -> str:
+        return self._btf_print
+
+    def __str__(self) -> str:
+        return self._btf_print
+
+    def _update_btf_print(self, title='easyidp.Container'):
+        key_list = list(self.item_label.keys())
+        num = len(key_list)
+        out_str = f'<{title}> with {num} items\n'
+        if num > 5:
+            for i, k in enumerate(key_list[:2]):
+                out_str += f"[{i}]\t{k}\n"
+                out_str += repr(self.id_item[self.item_label[k]])
+                out_str += '\n'
+            out_str += '...\n'
+            for i, k in enumerate(key_list[-2:]):
+                out_str += f"[{num-2+i}]\t{k}\n"
+                out_str += repr(self.id_item[self.item_label[k]])
+                out_str += '\n'
+        else:
+            for i, k in enumerate(key_list):
+                out_str += f"[{i}]\t{k}\n"
+                out_str += repr(self.id_item[self.item_label[k]])
+                out_str += '\n'
+
+        out_str = out_str[:-1]
+        self._btf_print = out_str
 
     def __len__(self):
         return len(self.id_item)
@@ -97,6 +128,8 @@ def __delitem__(self, key):
                 label = _find_key(self.item_label, idx)
                 self.item_label[label] = idx - 1
 
+        self._update_btf_print()
+
     def __iter__(self):
         return iter(self.id_item.values())
 

easyidp/roi.py

@@ -108,6 +108,8 @@ def read_shp(self, shp_path, shp_proj=None, name_field=None, include_title=False
         for k, v in roi_dict.items():
             self[k] = v
 
+        self._update_btf_print(title='easyidp.ROI')
+
     def read_labelme_json(self, json_path):
         """read roi from labelme marked json file
 
@@ -154,6 +156,8 @@ def read_labelme_json(self, json_path):
         else:
             raise TypeError(f"It seems [{json_path}] is not a Labelme json file.")
 
+        self._update_btf_print(title='easyidp.ROI')
+
     def change_crs(self, target_crs):
         """Change the geo coordinates of roi to another crs.
 

tests/test_init_class_func.py

@@ -1,6 +1,7 @@
 import re
 import sys
 import pytest
+import numpy as np
 import easyidp as idp
 
 def test_class_container():
@@ -50,7 +51,24 @@ def test_class_container():
     assert len(slice_test) == 3
     for k in slice_test.keys():
         assert k in ['6', '7', '8']
-
+
+def test_class_container_btf_print():
+    # short container
+    expected_str_s = '<easyidp.Container> with 3 items\n[0]\t1\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])\n[1]\t2\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])\n[2]\t3\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])'
+    ctn = idp.Container()
+    ctn['1'] = np.ones((4,3))
+    ctn['2'] = np.ones((4,3))
+    ctn['3'] = np.ones((4,3))
+
+    assert ctn._btf_print.replace(' ', '') ==  expected_str_s.replace(' ', '')
+
+    # long container
+    expected_str_l = '<easyidp.Container> with 6 items\n[0]\t1\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])\n[1]\t2\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])\n...\n[4]\t134\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])\n[5]\t135\narray([[1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.],\n       [1., 1., 1.]])'
+    ctn['133'] = np.ones((4,3))
+    ctn['134'] = np.ones((4,3))
+    ctn['135'] = np.ones((4,3))
+
+    assert ctn._btf_print.replace(' ', '') ==  expected_str_l.replace(' ', '')
 
 
 def test_def_parse_photo_relative_path():