From fba30e07d14cd645381a38e69822ac5c25cbe857 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sun, 30 Oct 2022 18:09:04 -0700
Subject: [PATCH] factor out mock python repl (#43)

---
 langchain/chains/python.py      |  4 +++-
 langchain/python.py             | 15 +++++++++++++++
 tests/unit_tests/test_python.py | 34 +++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 langchain/python.py
 create mode 100644 tests/unit_tests/test_python.py

diff --git a/langchain/chains/python.py b/langchain/chains/python.py
index 8d9e0e05dc..32311daed9 100644
--- a/langchain/chains/python.py
+++ b/langchain/chains/python.py
@@ -9,6 +9,7 @@ from typing import Dict, List
 from pydantic import BaseModel
 
 from langchain.chains.base import Chain
+from langchain.python import PythonREPL
 
 
 class PythonChain(Chain, BaseModel):
@@ -41,9 +42,10 @@ class PythonChain(Chain, BaseModel):
         return [self.output_key]
 
     def _run(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        python_repl = PythonREPL()
         old_stdout = sys.stdout
         sys.stdout = mystdout = StringIO()
-        exec(inputs[self.input_key])
+        python_repl.run(inputs[self.input_key])
         sys.stdout = old_stdout
         output = mystdout.getvalue()
         return {self.output_key: output}
diff --git a/langchain/python.py b/langchain/python.py
new file mode 100644
index 0000000000..b4a40cff17
--- /dev/null
+++ b/langchain/python.py
@@ -0,0 +1,15 @@
+"""Mock Python REPL."""
+from typing import Dict, Optional
+
+
+class PythonREPL:
+    """Simulates a standalone Python REPL."""
+
+    def __init__(self, _globals: Optional[Dict] = None, _locals: Optional[Dict] = None):
+        """Initialize with optional globals and locals."""
+        self._globals = _globals if _globals is not None else {}
+        self._locals = _locals if _locals is not None else {}
+
+    def run(self, command: str) -> None:
+        """Run command with own globals/locals."""
+        exec(command, self._globals, self._locals)
diff --git a/tests/unit_tests/test_python.py b/tests/unit_tests/test_python.py
new file mode 100644
index 0000000000..419f13cae3
--- /dev/null
+++ b/tests/unit_tests/test_python.py
@@ -0,0 +1,34 @@
+"""Test functionality of Python REPL."""
+
+import pytest
+
+from langchain.python import PythonREPL
+
+
+def test_python_repl() -> None:
+    """Test functionality when globals/locals are not provided."""
+    repl = PythonREPL()
+
+    # Run a simple initial command.
+    repl.run("foo = 1")
+    assert repl._locals["foo"] == 1
+
+    # Now run a command that accesses `foo` to make sure it still has it.
+    repl.run("bar = foo * 2")
+    assert repl._locals["bar"] == 2
+
+
+def test_python_repl_no_previous_variables() -> None:
+    """Test that it does not have access to variables created outside the scope."""
+    foo = 3  # noqa: F841
+    repl = PythonREPL()
+    with pytest.raises(NameError):
+        repl.run("print(foo)")
+
+
+def test_python_repl_pass_in_locals() -> None:
+    """Test functionality when passing in locals."""
+    _locals = {"foo": 4}
+    repl = PythonREPL(_locals=_locals)
+    repl.run("bar = foo * 2")
+    assert repl._locals["bar"] == 8