From 423b3b648e3d385c9d9c24009f0eefef7a50cb65 Mon Sep 17 00:00:00 2001
From: Victor Lin <13424970+victorlin@users.noreply.github.com>
Date: Fri, 9 Feb 2024 15:41:56 -0800
Subject: [PATCH] Support bool type in --query-columns

---
 augur/filter/include_exclude_rules.py               |  5 +++++
 augur/filter/io.py                                  |  2 +-
 tests/functional/filter/cram/filter-query-columns.t | 10 ++++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/augur/filter/include_exclude_rules.py b/augur/filter/include_exclude_rules.py
index f1c5bace2..a333914f7 100644
--- 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
index 56c113bb4..6ebf253ef 100644
--- 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
index 522e8eed6..807456d22 100644
--- 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]