wip: implemented len and getitem on iterators

Author: PJEstrada

sdk/diffgram/core/diffgram_dataset_iterator.py

@@ -21,6 +21,14 @@ def __iter__(self):
         self.current_file_index = 0
         return self
 
+    def __len__(self):
+        return len(self.diffgram_file_id_list)
+
+    def __getitem__(self, idx):
+        diffgram_file = self.project.file.get_by_id(self.diffgram_file_id_list[idx], with_instances = True)
+        instance_data = self.get_file_instances(diffgram_file)
+        return instance_data
+
     def __next__(self):
         file_id = self.diffgram_file_id_list[self.current_file_index]
         diffgram_file = self.project.file.get_by_id(file_id, with_instances = True)

sdk/diffgram/core/directory.py

@@ -3,6 +3,8 @@
 import logging
 from diffgram.pytorch_diffgram.diffgram_pytorch_dataset import DiffgramPytorchDataset
 from diffgram.tensorflow_diffgram.diffgram_tensorflow_dataset import DiffgramTensorflowDataset
+from diffgram.core.diffgram_dataset_iterator import DiffgramDatasetIterator
+
 
 def get_directory_list(self):
 	"""
@@ -70,14 +72,19 @@ def set_directory_by_name(self, name):
 				  str(names_attempted))
 
 
-class Directory():
+class Directory(DiffgramDatasetIterator):
 
-	def __init__(self,
-			     client):
+	def __init__(self, client, file_id_list_sliced = None):
 
 		self.client = client
 		self.id = None
 		self.file_list_metadata = {}
+		
+		if file_id_list_sliced is None:
+			self.file_id_list = self.all_file_ids()
+		else:
+			self.file_id_list = file_id_list_sliced
+		super(Directory, self).__init__(self.client, self.file_id_list)
 
 	def all_files(self):
 		"""

sdk/diffgram/core/sliced_directory.py

@@ -11,6 +11,8 @@ def __init__(self, client, original_directory: Directory, query: str):
         self.client = client
         # Share the same ID from the original directory as this is just an in-memory construct for better semantics.
         self.id = original_directory.id
+        self.file_id_list = self.all_file_ids()
+        super(Directory, self).__init__(self.client, self.file_id_list)
 
     def all_file_ids(self):
         page_num = 1
@@ -29,10 +31,10 @@ def to_pytorch(self, transform = None):
             Transforms the file list inside the dataset into a pytorch dataset.
         :return:
         """
-        file_id_list = self.all_file_ids()
+
         pytorch_dataset = DiffgramPytorchDataset(
             project = self.client,
-            diffgram_file_id_list = file_id_list,
+            diffgram_file_id_list = self.file_id_list,
             transform = transform
 
         )