Fix register kwarg in chained expected datasets (issue #972)

Author: k1o0

ibllib/oneibl/data_handlers.py

@@ -143,6 +143,7 @@ def find_files(self, session_path, register=False):
         - Currently if `unique` is true and multiple files are found, all files are returned without an exception raised
           although this may change in the future.
         - If `register` is false, all files are returned regardless of whether they are intended to be registered.
+        - If `register` is true, an input with register=True may not be returned if part of an OR operation.
         - If `inverted` is true, and files are found, the glob pattern is returned as missing.
         - If XOR, returns all patterns if all are present when only one should be, otherwise returns all missing
           patterns.
@@ -163,22 +164,22 @@ def find_files(self, session_path, register=False):
                 missing = self.glob_pattern
         elif self.operator == 'and':
             assert len(self._identifiers) == 2
-            _ok, _actual_files, _missing = zip(*map(lambda x: x.find_files(session_path), self._identifiers))
+            _ok, _actual_files, _missing = zip(*map(lambda x: x.find_files(session_path, register=register), self._identifiers))
             ok = all(_ok)
             actual_files = flatten(_actual_files)
             missing = set(filter(None, flatten(_missing)))
         elif self.operator == 'or':
             assert len(self._identifiers) == 2
             missing = set()
             for d in self._identifiers:
-                ok, actual_files, _missing = d.find_files(session_path)
+                ok, actual_files, _missing = d.find_files(session_path, register=register)
                 if ok:
                     break
                 if missing is not None:
                     missing.update(_missing) if isinstance(_missing, set) else missing.add(_missing)
         elif self.operator == 'xor':
             assert len(self._identifiers) == 2
-            _ok, _actual_files, _missing = zip(*map(lambda x: x.find_files(session_path), self._identifiers))
+            _ok, _actual_files, _missing = zip(*map(lambda x: x.find_files(session_path, register=register), self._identifiers))
             ok = sum(_ok) == 1  # and sum(map(bool, map(len, _actual_files))) == 1
             # Return only those datasets that are complete if OK
             actual_files = _actual_files[_ok.index(True)] if ok else flatten(_actual_files)

ibllib/oneibl/registration.py

@@ -16,7 +16,6 @@
 from iblutil.util import ensure_list
 
 import ibllib
-import ibllib.io.extractors.base
 from ibllib.time import isostr2date
 import ibllib.io.raw_data_loaders as raw
 from ibllib.io import session_params

ibllib/tests/test_oneibl.py

@@ -854,6 +854,7 @@ def setUp(self):
         self.img_path.joinpath('imaging.frames.tar.bz2').touch()
 
     def test_and(self):
+        """Test for ExpectedDataset input with AND operator."""
         I = handlers.ExpectedDataset.input  # noqa
         sig = I('*.tif', 'raw_imaging_data_[0-9]*') & I('imaging.frames.tar.bz2', 'raw_imaging_data_[0-9]*')
         self.assertEqual('and', sig.operator)
@@ -863,6 +864,16 @@ def test_and(self):
         self.assertEqual('foo_00000.tif', files[0].name)
         self.assertEqual('imaging.frames.tar.bz2', files[-1].name)
         self.assertEqual(set(), missing)
+        # Check with register arg
+        ok, _files, missing = sig.find_files(self.session_path, register=True)
+        self.assertTrue(ok)
+        self.assertEqual([], _files)
+        self.assertEqual(set(), missing)
+        sig._identifiers[1].register = True
+        ok, _files, missing = sig.find_files(self.session_path, register=True)
+        self.assertTrue(ok)
+        self.assertEqual(1, len(_files))
+        self.assertEqual('imaging.frames.tar.bz2', _files[0].name)
         # Deleting one tif file shouldn't affect the signature
         files[0].unlink()
         ok, files, missing = sig.find_files(self.session_path)
@@ -878,6 +889,7 @@ def test_and(self):
         self.assertEqual({'raw_imaging_data_[0-9]*/imaging.frames.tar.bz2'}, missing)
 
     def test_or(self):
+        """Test for ExpectedDataset input with OR operator."""
         I = handlers.ExpectedDataset.input  # noqa
         sig = I('*.tif', 'raw_imaging_data_[0-9]*') | I('imaging.frames.tar.bz2', 'raw_imaging_data_[0-9]*')
         self.assertEqual('or', sig.operator)
@@ -886,6 +898,17 @@ def test_or(self):
         self.assertEqual(5, len(files))
         self.assertEqual({'.tif'}, set(x.suffix for x in files))
         self.assertEqual(set(), missing)
+        # Test with register filter
+        ok, _files, missing = sig.find_files(self.session_path, register=True)
+        self.assertTrue(ok)
+        self.assertEqual([], _files)
+        self.assertEqual(set(), missing)
+        sig._identifiers[1].register = True
+        ok, _files, missing = sig.find_files(self.session_path, register=True)
+        self.assertTrue(ok)
+        # NB: Although 2nd file has register=True, the first is present so nothing is returned here
+        self.assertEqual([], _files)
+        self.assertEqual(set(), missing)
         for f in files:
             f.unlink()
         ok, files, missing = sig.find_files(self.session_path)
@@ -900,6 +923,7 @@ def test_or(self):
         self.assertEqual({'raw_imaging_data_[0-9]*/*.tif', 'raw_imaging_data_[0-9]*/imaging.frames.tar.bz2'}, missing)
 
     def test_xor(self):
+        """Test for ExpectedDataset input with XOR operator."""
         I = handlers.ExpectedDataset.input  # noqa
         sig = I('*.tif', 'raw_imaging_data_[0-9]*') ^ I('imaging.frames.tar.bz2', 'raw_imaging_data_[0-9]*')
         self.assertEqual('xor', sig.operator)
@@ -909,6 +933,17 @@ def test_xor(self):
         self.assertEqual({'.tif', '.bz2'}, set(x.suffix for x in files))
         expected_missing = {'raw_imaging_data_[0-9]*/*.tif', 'raw_imaging_data_[0-9]*/imaging.frames.tar.bz2'}
         self.assertEqual(expected_missing, missing)
+        # Test register kwarg
+        ok, _files, missing = sig.find_files(self.session_path, register=True)
+        self.assertFalse(ok)
+        self.assertEqual([], _files)
+        self.assertEqual(expected_missing, missing)
+        sig._identifiers[1].register = True
+        ok, _files, _ = sig.find_files(self.session_path, register=True)
+        self.assertFalse(ok)
+        self.assertEqual(1, len(_files))
+        self.assertEqual('imaging.frames.tar.bz2', _files[0].name)
+        # Remove first expected dataset(s)
         for f in filter(lambda x: x.suffix == '.tif', files):
             f.unlink()
         ok, files, missing = sig.find_files(self.session_path)