From 53b7da5c75e94dd7d5a49520fb11a26df18fe16a Mon Sep 17 00:00:00 2001 From: Christophe Bedard Date: Sun, 10 Nov 2019 16:05:23 -0800 Subject: [PATCH] Allow single string input for time columns --- tracetools_analysis/tracetools_analysis/utils.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tracetools_analysis/tracetools_analysis/utils.py b/tracetools_analysis/tracetools_analysis/utils.py index f319cf7..d977900 100644 --- a/tracetools_analysis/tracetools_analysis/utils.py +++ b/tracetools_analysis/tracetools_analysis/utils.py @@ -53,19 +53,24 @@ class DataModelUtil(): @staticmethod def convert_time_columns( original: DataFrame, - columns_ns_to_ms: List[str] = [], - columns_ns_to_datetime: List[str] = [], + columns_ns_to_ms: Union[List[str], str] = [], + columns_ns_to_datetime: Union[List[str], str] = [], inplace: bool = True, ) -> DataFrame: """ Convert time columns from nanoseconds to either milliseconds or `datetime` objects. :param original: the original `DataFrame` - :param columns_ns_to_ms: the columns for which to convert ns to ms - :param columns_ns_to_datetime: the columns for which to convert ns to `datetime` + :param columns_ns_to_ms: the column(s) for which to convert ns to ms + :param columns_ns_to_datetime: the column(s) for which to convert ns to `datetime` :param inplace: whether to convert in place or to return a copy :return: the resulting `DataFrame` """ + if not isinstance(columns_ns_to_ms, list): + columns_ns_to_ms = list(columns_ns_to_ms) + if not isinstance(columns_ns_to_datetime, list): + columns_ns_to_datetime = list(columns_ns_to_datetime) + df = original if inplace else original.copy() # Convert from ns to ms if len(columns_ns_to_ms) > 0: