Remove

Author: benedekrozemberczki

chemicalx/data/contextfeatureset.py

@@ -134,15 +134,15 @@ def get_context_feature_count(self) -> int:
         """Getting the number of feature dimensions.
 
         Returns:
-            int: The number of feature dimensions.
+            feature_count (int): The number of feature dimensions.
         """
+        feature_count = None
         if len(self.__dict__) > 0:
             contexts = list(self.keys())
             first_context = contexts[0]
             feature_vector = self.__dict__[first_context]
-            return feature_vector.shape[1]
-        else:
-            return 0
+            feature_count = feature_vector.shape[1]
+        return feature_count
 
     def get_feature_matrix(self, contexts: List[str]) -> np.ndarray:
         """Getting the feature matrix for a list of contexts.

tests/unit/test_datastructures.py

@@ -16,11 +16,16 @@ def setUp(self):
     def test_get(self):
         assert self.context_feature_set["context_2"].shape == (1, 3)
         assert ("context_2" in self.context_feature_set) == True
+        assert self.context_feature_set.has_context("context_2") == True
 
     def test_copy(self):
         self.another_context_feature_set = self.context_feature_set.copy()
         assert len(self.another_context_feature_set) == len(self.context_feature_set)
 
+        del self.another_context_feature_set["context_1"]
+        del self.another_context_feature_set["context_2"]
+        assert self.another_context_feature_set.get_context_feature_count() == None
+
     def test_len(self):
         assert len(self.context_feature_set) == 2
 
@@ -48,6 +53,10 @@ def test_iteration(self):
             feature_vector = self.context_feature_set[context]
             assert feature_vector.shape == (1, 3)
 
+    def test_clearing(self):
+        self.context_feature_set.clear()
+        assert len(self.context_feature_set) == 0
+
 
 class TestDrugFeatureSet(unittest.TestCase):
     """