pyp: __pyp_before__ (#39)

Author: hauntsaninja

pyp.py

@@ -313,8 +313,17 @@ def parse_input(code: List[str]) -> ast.Module:
                 raise PypError("".join(message).strip()) from e
 
         self.before_tree = parse_input(before)
+        if "__pyp_before__" in config.name_to_def:
+            config_before = config.parts[config.name_to_def["__pyp_before__"]]
+            if not isinstance(config_before, ast.FunctionDef):
+                raise PypError("Config __pyp_before__ must be a function")
+            self.before_tree.body = config_before.body + self.before_tree.body
+
         self.tree = parse_input(code)
+
         self.after_tree = parse_input(after)
+        if "__pyp_after__" in config.name_to_def:
+            raise PypError("Config __pyp_after__ not supported")
 
         f = NameFinder(self.before_tree, self.tree, self.after_tree)
         self.defined: Set[str] = f.top_level_defined

tests/test_pyp.py

@@ -697,6 +697,26 @@ def test_config_conditional_current_shortcoming(config_mock):
     compare_scripts(run_pyp(["--explain", "unparse(ast.parse('x')); pass"]), script3)
 
 
+@patch("pyp.get_config_contents")
+def test_config_pyp_before(config_mock):
+    config_mock.return_value = """
+import signal
+import sys
+def __pyp_before__():
+    signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+"""
+    script = r"""#!/usr/bin/env python3
+import signal
+import sys
+signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+for x in sys.stdin:
+    x = x.rstrip('\n')
+    if x is not None:
+        print(x)
+"""
+    compare_scripts(run_pyp(["--explain", "x"]), script)
+
+
 def test_config_end_to_end(monkeypatch, capsys):
     with tempfile.NamedTemporaryFile("w") as f:
         config = "def foo(): return 1"