Merge pull request #1119 from effigies/fix/padded-int-arrays

fix: Allow numpy arrays to treat PaddedInt like int

Author: effigies

src/bids/layout/tests/test_utils.py

@@ -2,12 +2,14 @@
 
 import os
 
+import numpy as np
 import pytest
+import pandas as pd
 import bids
 from bids.exceptions import ConfigError
 
 from ..models import Entity, Config
-from ..utils import BIDSMetadata, parse_file_entities, add_config_paths
+from ..utils import BIDSMetadata, PaddedInt, parse_file_entities, add_config_paths
 
 
 def test_bidsmetadata_class():
@@ -73,3 +75,31 @@ def test_add_config_paths():
     add_config_paths(dummy=bids_json)
     config = Config.load('dummy')
     assert 'subject' in config.entities
+
+
+def test_PaddedInt_array_comparisons():
+    # Array comparisons should work, not raise exceptions
+    arr = np.array([5, 5, 5])
+    assert np.all(arr == PaddedInt(5))
+    assert np.all(arr == PaddedInt("05"))
+    assert np.all(PaddedInt(5) == arr)
+    assert np.all(PaddedInt("05") == arr)
+
+    # If the value gets put into an array, it should be considered an int
+    # We lose the padding, but it's unlikely we would try to recover it
+    # from an array.
+    assert np.array([PaddedInt(5)]).dtype == np.int64
+
+    # Verify that we do get some False results
+    assert np.array_equal(np.array([4, 5, 6]) == PaddedInt(5), [False, True, False])
+
+
+def test_PaddedInt_dataframe_behavior():
+    # Verify that pandas dataframes are not more special than numpy arrays,
+    # as far as PaddedInt is concerned
+    df = pd.DataFrame({'a': [5, 5, 5]})
+    pidf = pd.DataFrame({'a': [PaddedInt(5)] * 3})
+    assert np.all(df['a'] == PaddedInt(5))
+    assert np.all(df == pidf)
+
+    assert pidf['a'].dtype is np.dtype('int64')

src/bids/layout/utils.py

@@ -65,7 +65,13 @@ def __init__(self, val):
         self.sval = str(val)
 
     def __eq__(self, val):
-        return val == self.sval or super().__eq__(val)
+        try:
+            return val == self.sval or super().__eq__(val)
+        except ValueError:
+            # `or` triggers `__bool__`. If this fails, it is almost
+            # certainly an array type. Do not attempt string comparisons
+            # and allow val to determine the return type
+            return val == self
 
     def __str__(self):
         return self.sval