Support bool type in --query-columns

nextstrain · Feb 9, 2024 · 423b3b6 · 423b3b6
1 parent bd35393
commit 423b3b6
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/augur/filter/include_exclude_rules.py b/augur/filter/include_exclude_rules.py
@@ -232,6 +232,11 @@ def filter_by_query(metadata: pd.DataFrame, query: str, column_types: Optional[D
                 metadata_copy[column] = pd.to_numeric(metadata_copy[column], errors='raise', downcast='float')
             except ValueError as e:
                 raise AugurError(f"Failed to convert value in column {column!r} to float. {e}")
+        elif dtype == 'bool':
+            try:
+                metadata_copy[column] = metadata_copy[column].map(_string_to_boolean)
+            except ValueError as e:
+                raise AugurError(f"Failed to convert value in column {column!r} to bool. {e}")
         elif dtype == 'str':
             metadata_copy[column] = metadata_copy[column].astype('str', errors='ignore')
 

diff --git a/augur/filter/io.py b/augur/filter/io.py
@@ -128,7 +128,7 @@ def write_metadata_based_outputs(input_metadata_path: str, delimiters: Sequence[
 
 
 # These are the types accepted in the following function.
-ACCEPTED_TYPES = {'int', 'float', 'str'}
+ACCEPTED_TYPES = {'int', 'float', 'bool', 'str'}
 
 def column_type_pair(input: str):
     """Get a 2-tuple for column name to type.

diff --git a/tests/functional/filter/cram/filter-query-columns.t b/tests/functional/filter/cram/filter-query-columns.t
@@ -53,3 +53,13 @@ Specifying category:float does not work.
   >  --output-strains filtered_strains.txt
   ERROR: Failed to convert value in column 'category' to float. Unable to parse string "A" at position 0
   [2]
+
+Specifying category:bool also does not work.
+
+  $ ${AUGUR} filter \
+  >  --metadata metadata.tsv \
+  >  --query "coverage >= 0.95 & category == 'B'" \
+  >  --query-columns category:bool \
+  >  --output-strains filtered_strains.txt
+  ERROR: Failed to convert value in column 'category' to bool. Unable to convert 'A' to a boolean value.
+  [2]