Skip to content

Commit 708ce8b

Browse files
authored
pyp: __pyp_before__ (#39)
1 parent 5af2a58 commit 708ce8b

File tree

2 files changed

+29
-0
lines changed

2 files changed

+29
-0
lines changed

pyp.py

+9
Original file line numberDiff line numberDiff line change
@@ -313,8 +313,17 @@ def parse_input(code: List[str]) -> ast.Module:
313313
raise PypError("".join(message).strip()) from e
314314

315315
self.before_tree = parse_input(before)
316+
if "__pyp_before__" in config.name_to_def:
317+
config_before = config.parts[config.name_to_def["__pyp_before__"]]
318+
if not isinstance(config_before, ast.FunctionDef):
319+
raise PypError("Config __pyp_before__ must be a function")
320+
self.before_tree.body = config_before.body + self.before_tree.body
321+
316322
self.tree = parse_input(code)
323+
317324
self.after_tree = parse_input(after)
325+
if "__pyp_after__" in config.name_to_def:
326+
raise PypError("Config __pyp_after__ not supported")
318327

319328
f = NameFinder(self.before_tree, self.tree, self.after_tree)
320329
self.defined: Set[str] = f.top_level_defined

tests/test_pyp.py

+20
Original file line numberDiff line numberDiff line change
@@ -697,6 +697,26 @@ def test_config_conditional_current_shortcoming(config_mock):
697697
compare_scripts(run_pyp(["--explain", "unparse(ast.parse('x')); pass"]), script3)
698698

699699

700+
@patch("pyp.get_config_contents")
701+
def test_config_pyp_before(config_mock):
702+
config_mock.return_value = """
703+
import signal
704+
import sys
705+
def __pyp_before__():
706+
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
707+
"""
708+
script = r"""#!/usr/bin/env python3
709+
import signal
710+
import sys
711+
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
712+
for x in sys.stdin:
713+
x = x.rstrip('\n')
714+
if x is not None:
715+
print(x)
716+
"""
717+
compare_scripts(run_pyp(["--explain", "x"]), script)
718+
719+
700720
def test_config_end_to_end(monkeypatch, capsys):
701721
with tempfile.NamedTemporaryFile("w") as f:
702722
config = "def foo(): return 1"

0 commit comments

Comments
 (0)