Added where eq, neq for single column when satisfied

Author: kanhereashwin

gnss_lib_py/parsers/navdata.py

@@ -377,19 +377,21 @@ def argwhere(self, key_idx, value, condition="eq"):
                                + "for string condition checks")
             # Extract columns where condition holds true and return new NavData
             if condition == "eq":
-                new_cols = np.argwhere(np.isin(self[row, :],str_check))
+                new_cols = np.argwhere(np.atleast_1d(np.isin(self[row, :],
+                                                             str_check)))
             else:
                 # condition == "neq"
-                new_cols = np.argwhere(~np.isin(self[row, :],str_check))
+                new_cols = np.argwhere(np.atleast_1d(~np.isin(self[row, :],
+                                                              str_check)))
 
         else:
             # Values in row are numerical
             # Find columns where value can be found and return new NavData
             if condition=="eq":
                 if isinstance(value,(np.ndarray,list,tuple,set)):
                     # use numpy's isin() condition if list of values
-                    new_cols = np.argwhere(np.isin(self.array[row, :],
-                                           value))
+                    new_cols = np.argwhere(np.atleast_1d(np.isin(self.array[row, :],
+                                           value)))
                 elif not isinstance(value,str) and np.isnan(value):
                     # check isinstance b/c np.isnan can't handle strings
                     new_cols = np.argwhere(np.isnan(self.array[row, :]))
@@ -398,11 +400,11 @@ def argwhere(self, key_idx, value, condition="eq"):
             elif condition=="neq":
                 if isinstance(value,(np.ndarray,list,tuple,set)):
                     # use numpy's isin() condition if list of values
-                    new_cols = np.argwhere(~np.isin(self.array[row, :],
-                                           value))
+                    new_cols = np.argwhere(np.atleast_1d(~np.isin(self.array[row, :],
+                                           value)))
                 elif not isinstance(value,str) and np.isnan(value):
                     # check isinstance b/c np.isnan can't handle strings
-                    new_cols = np.argwhere(~np.isnan(self.array[row, :]))
+                    new_cols = np.argwhere(np.atleast_1d(~np.isnan(self.array[row, :])))
                 else:
                     new_cols = np.argwhere(self.array[row, :]!=value)
             elif condition == "leq":

tests/parsers/test_navdata.py

@@ -1690,6 +1690,19 @@ def test_where_str(csv_simple):
     compare_df = compare_df[compare_df['strings']!="gps"].reset_index(drop=True)
     pd.testing.assert_frame_equal(data_small.pandas_df(), compare_df)
 
+    #Test equality for cases where there is only one column
+    data_single_column = data.where('strings', 'glonass', 'eq')
+    data_new = data_single_column.where('strings', 'glonass', 'eq')
+    compare_df = data.pandas_df()
+    compare_df = compare_df[compare_df['strings']=="glonass"].reset_index(drop=True)
+    pd.testing.assert_frame_equal(data_new.pandas_df(), compare_df)
+
+    #Test inequality for cases where there is only one column
+    data_new = data_single_column.where('strings', 'gps', 'neq')
+    # Both cases should return the same dataframe as before
+    pd.testing.assert_frame_equal(data_new.pandas_df(), compare_df)
+
+
 def test_where_empty(df_simple):
     """Verify empty slices.
 
@@ -1750,6 +1763,17 @@ def test_where_numbers(csv_simple):
         compare_df = compare_df.iloc[pd_rows[idx], :].reset_index(drop=True)
         pd.testing.assert_frame_equal(data_small.pandas_df(), compare_df)
 
+    #Test equality for cases where there is only one column
+    data_single_column = data.where('integers', 10, 'eq')
+    data_new = data_single_column.where('integers', 10, 'eq')
+    compare_df = data.pandas_df()
+    compare_df = compare_df[compare_df['integers']==10].reset_index(drop=True)
+    pd.testing.assert_frame_equal(data_new.pandas_df(), compare_df)
+
+    #Test inequality for cases where there is only one column
+    data_new = data_single_column.where('integers', 56, 'neq')
+    pd.testing.assert_frame_equal(data_new.pandas_df(), compare_df)
+
 def test_where_errors(csv_simple):
     """Testing error cases for NavData.where