Add chunk getting with integer values (#55)

* Add chunk getting with integer values

* Add field getting for chunk json

* Update test_doc_utils.py

* Update test_doc_utils.py

* Update test_doc_utils.py

Author: boba-and-beer

tests/test_doc_utils.py

@@ -1,11 +1,25 @@
 """Testing for document utilities.
 """
+import pytest
+from vectorai.errors import MissingFieldError
 
 def test_set_field(test_client):
     sample = {}
     test_client.set_field("simple", doc=sample, value=[0, 2])
     assert test_client.get_field("simple", sample) == [0, 2]
 
+def test_get_field_chunk(test_client):
+    sample = {
+        'kfc': [{'food': 'chicken'}, {'food': 'prawns'}]}
+    assert test_client.get_field('kfc.0.food', sample) == 'chicken'
+    assert test_client.get_field('kfc.1.food', sample) == 'prawns'
+
+def test_get_field_chunk_error(test_client):
+    sample = {
+        'kfc': [{'food': 'chicken'}, {'food': 'prawns'}]}
+    with pytest.raises(MissingFieldError):
+        test_client.get_field('kfc.food', sample)
+
 def test_get_fields(test_client):
     doc = test_client.create_sample_documents(1)[0]
     assert len(test_client.get_fields(['size.cm', 'size.feet'], doc)) == 2

vectorai/doc_utils.py

@@ -33,12 +33,29 @@ def get_field(self, field: str, doc: Dict):
         for f in field.split("."):
             try:
                 d = d[f]
-            except:
+            except KeyError:
                 try:
                     return doc[field]
-                except:
+                except KeyError:
                     raise MissingFieldError("Document is missing " + field)
+            except TypeError:
+                if self._is_string_integer(f):
+                    # Get the Get the chunk document out.
+                    d = d[int(f)]
+                else:
+                    raise MissingFieldError("Document is missing " + f + ' of ' + field)
         return d
+    
+    @classmethod
+    def _is_string_integer(cls, x):
+        """Test if a string is numeric
+        """
+        try:
+            int(x)
+            return True
+        except:
+            return False
+
 
     @classmethod
     def get_fields(self, fields: List[str], doc: Dict) -> List[Any]:

vectorai/utils.py

@@ -348,6 +348,15 @@ def show_json(self, json: dict, selected_fields: List[str]=None, image_fields: L
             image_fields=image_fields, audio_fields=audio_fields, image_width=image_width, include_vector=include_vector)
         return self.show_df(self.results_to_df(json).head(nrows)[image_fields + audio_fields + selected_fields], 
             image_fields=image_fields, audio_fields=audio_fields, image_width=image_width, include_vector=include_vector)
+    
+    def show_chunk_json(self, json: dict, selected_fields: List[str]=None, image_fields: List[str]=[], 
+    audio_fields: List[str]=[], nrows: int=5, image_width: int=60, include_vector=False):
+        """Show results if the documents are chunked.
+        For images, concatenates the chunk images into the same numpy array
+        For text, puts them one after the other with smaller index.
+        No Audio chunking for now. 
+        """
+        raise NotImplementedError
 
 def get_random_int(low=0, high=9999):
     """