diff --git a/orca/orca.py b/orca/orca.py index f02c909..38ca7fc 100644 --- a/orca/orca.py +++ b/orca/orca.py @@ -1911,6 +1911,30 @@ def injectables(**kwargs): _INJECTABLES = original +@contextmanager +def tables(**kwargs): + """ + Temporarily set DataFrames as registered tables. + + Tables will be returned to their original state when the context + manager exits. Caching is not enabled for tables registered via + this function. + + """ + global _TABLES + + original = _TABLES.copy() + + for k, v in kwargs.items(): + if not isinstance(v, pd.DataFrame): + raise ValueError('tables only accepts DataFrames') + add_table(k, v) + + yield + + _TABLES = original + + def eval_variable(name, **kwargs): """ Execute a single variable function registered with Orca diff --git a/orca/tests/test_orca.py b/orca/tests/test_orca.py index a19d7bd..d290452 100644 --- a/orca/tests/test_orca.py +++ b/orca/tests/test_orca.py @@ -1024,6 +1024,18 @@ def test_injectables_cm(): } +def test_tables_cm(): + orca.add_table('a', pd.DataFrame()) + + with orca.tables(): + assert sorted(orca._TABLES.keys()) == ['a'] + + with orca.tables(a=pd.DataFrame(), b=pd.DataFrame()): + assert sorted(orca._TABLES.keys()) == ['a', 'b'] + + assert sorted(orca._TABLES.keys()) == ['a'] + + def test_is_expression(): assert orca.is_expression('name') is False assert orca.is_expression('table.column') is True