context manager for temporarily registering tables

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

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