diff --git a/dftimewolf/lib/utils.py b/dftimewolf/lib/utils.py index 596f4a3f9..e25dddd97 100644 --- a/dftimewolf/lib/utils.py +++ b/dftimewolf/lib/utils.py @@ -3,6 +3,7 @@ import abc import argparse +import datetime import os import random import re @@ -102,6 +103,10 @@ def ImportArgsFromDict(value: Any, It is used to load arguments from the CLI and any extra configuration parameters passed in recipes. + When a token resolves to a ``datetime.datetime`` and the recipe value is + solely ``@token``, the datetime object is preserved. When the token is + embedded in a larger string, it is formatted as an ISO8601 timestamp. + Args: value (object): The value a dictionary. This is passed recursively and may change in nature: string, list, or dict. The top-level variable should @@ -120,6 +125,14 @@ def ImportArgsFromDict(value: Any, actual_param = args[token] if isinstance(actual_param, str): value = value.replace("@"+token, actual_param) + elif isinstance(actual_param, datetime.datetime): + # Keep datetime objects when the recipe arg is solely "@token" so + # modules that expect datetime continue to receive them. When the + # token is embedded in a larger string, expand to ISO8601. + if value == f'@{token}': + value = actual_param + else: + value = value.replace('@' + token, actual_param.isoformat()) else: value = actual_param elif isinstance(value, list): diff --git a/tests/test_dftimewolf.py b/tests/test_dftimewolf.py index b9f035e87..0a2e1257d 100644 --- a/tests/test_dftimewolf.py +++ b/tests/test_dftimewolf.py @@ -2,6 +2,7 @@ """Tests for DFTimewolf functions.""" import argparse +import datetime import unittest import six @@ -139,3 +140,37 @@ def test_config_fills_missing_args(self): imported_args = dftw_utils.ImportArgsFromDict( provided_args, vars(args), config.Config) self.assertEqual(imported_args, expected_args) + + def test_datetime_token_preserves_object(self): + """Tests that a sole @datetime token remains a datetime object.""" + start_date = datetime.datetime( + 2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc) + recipe_args = {'start_time': '@start_date'} + imported_args = dftw_utils.ImportArgsFromDict( + recipe_args, {'start_date': start_date}, config.Config) + + self.assertIs(imported_args['start_time'], start_date) + + def test_datetime_tokens_expand_to_iso8601_in_strings(self): + """Tests that datetime tokens embedded in strings expand to ISO8601.""" + start_date = datetime.datetime( + 2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc) + end_date = datetime.datetime( + 2024, 1, 2, 18, 30, 45, tzinfo=datetime.timezone.utc) + recipe_args = { + 'sheet_title': 'example @start_date @end_date', + } + expected_args = { + 'sheet_title': ( + f'example {start_date.isoformat()} {end_date.isoformat()}'), + } + + imported_args = dftw_utils.ImportArgsFromDict( + recipe_args, + {'start_date': start_date, 'end_date': end_date}, + config.Config) + + self.assertEqual(imported_args, expected_args) + self.assertEqual( + imported_args['sheet_title'], + 'example 2024-01-01T12:00:00+00:00 2024-01-02T18:30:45+00:00')