Prevent custom entities from drawing standard ent

A failing test drew "ce" be chance, which is a valid entity name
khanlab · Feb 21, 2024 · 6161825 · 6161825
1 parent 35cc8e0
commit 6161825
Showing 1 changed file with 23 additions and 21 deletions.
diff --git a/snakebids/tests/test_paths/test_bids.py b/snakebids/tests/test_paths/test_bids.py
@@ -44,29 +44,31 @@ def _bids_args(
         nonstandard: bool = True,
         custom: bool = True,
     ):
+        std_ents = (
+            sb_st.bids_entity(whitelist_entities=entities)
+            if entities is not None
+            else sb_st.nothing()
+        )
+        custom_ents = (
+            _values()
+            .map(BidsEntity)
+            .filter(
+                lambda s: str(s)
+                not in (
+                    {"datatype", "suffix", "extension", "prefix"} | set(std_entities)
+                )
+            )
+            if custom
+            else sb_st.nothing()
+        )
+        nonstd_ents = (
+            sb_st.bids_entity(whitelist_entities=["datatype", "suffix", "extension"])
+            if nonstandard
+            else sb_st.nothing()
+        )
         return (
             st.dictionaries(
-                keys=st.one_of(
-                    # standard
-                    sb_st.bids_entity(whitelist_entities=entities)
-                    if entities is not None
-                    else sb_st.nothing(),
-                    # custom entities
-                    _values()
-                    .map(BidsEntity)
-                    .filter(
-                        lambda s: str(s)
-                        not in {"datatype", "suffix", "extension", "prefix"}
-                    )
-                    if custom
-                    else sb_st.nothing(),
-                    # nonstandard entities
-                    sb_st.bids_entity(
-                        whitelist_entities=["datatype", "suffix", "extension"]
-                    )
-                    if nonstandard
-                    else sb_st.nothing(),
-                ),
+                keys=std_ents | custom_ents | nonstd_ents,
                 # The boolean here is to decide whether to use the entity or the tag in
                 # the BidsEntity generated above
                 values=st.tuples(