From fc746b2a5f992b0f5538727c6402f6e6047ed822 Mon Sep 17 00:00:00 2001 From: Toby Mao Date: Thu, 30 Jan 2025 09:35:07 -0800 Subject: [PATCH] chore: improve load times by making import optional (#3748) --- sqlmesh/core/test/definition.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py index aba4f2080..c79385f4a 100644 --- a/sqlmesh/core/test/definition.py +++ b/sqlmesh/core/test/definition.py @@ -12,7 +12,6 @@ import numpy as np import pandas as pd -import time_machine from io import StringIO from pandas.api.types import is_object_dtype from sqlglot import Dialect, exp @@ -670,13 +669,15 @@ def runTest(self) -> None: def _execute_model(self) -> pd.DataFrame: """Executes the python model and returns a DataFrame.""" - time_ctx = ( - time_machine.travel(self._execution_time, tick=False) - if self._execution_time - else nullcontext() - ) + if self._execution_time: + import time_machine + + time_ctx: AbstractContextManager = time_machine.travel(self._execution_time, tick=False) + else: + time_ctx = nullcontext() + with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms): - with t.cast(AbstractContextManager, time_ctx): + with time_ctx: variables = self.body.get("vars", {}).copy() time_kwargs = { key: variables.pop(key) for key in TIME_KWARG_KEYS if key in variables