Skip to content

Commit 32b4e93

Browse files
committed
Allow to select col to search by col number on top of col name
1 parent 1036287 commit 32b4e93

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

test/unit/utilities.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -296,10 +296,15 @@ def csv_col_has_values(csv_file: str, col_name: str, *values) -> bool:
296296
return False
297297

298298

299-
def csv_col_count_values(csv_file: str, col_name: str, *values) -> int:
299+
def csv_col_count_values(csv_file: str, col_name_or_nbr: Union[str, int], *values) -> int:
300+
"""Counts the number of times a given column has one of the given values"""
300301
values_to_search = list(values).copy()
301302
with open(csv_file, encoding="utf-8") as fd:
302-
(col,) = get_cols(next(reader := csv.reader(fd)), col_name)
303+
reader = csv.reader(fd)
304+
if isinstance(col_name_or_nbr, int):
305+
col = col_name_or_nbr - 1
306+
else:
307+
(col,) = get_cols(next(reader, col_name_or_nbr))
303308
counter = sum(1 if line[col] in values_to_search else 0 for line in reader)
304309
return counter
305310

0 commit comments

Comments
 (0)