From bb1913109e7d4c9cf52259b62b3959d2a24b27f2 Mon Sep 17 00:00:00 2001 From: Josh Karpel Date: Wed, 19 May 2021 21:29:54 -0500 Subject: [PATCH] Improve `get_function_body` (#38) --- spiel/deck.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/spiel/deck.py b/spiel/deck.py index 906dcb6..f2b66f3 100644 --- a/spiel/deck.py +++ b/spiel/deck.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dis import inspect import sys from collections.abc import Collection @@ -74,16 +75,10 @@ class Deck(Collection): def get_function_body(function: Callable) -> str: - lines, _ = inspect.getsourcelines(function) - - prev_indent = None - for idx, line in enumerate(lines): - if prev_indent is None: - prev_indent = count_leading_whitespace(line) - elif count_leading_whitespace(line) > prev_indent: - return dedent("".join(lines[idx:])) - - raise ValueError(f"Could not extract function body from {function}") + lines, line_of_def_start = inspect.getsourcelines(function) + line_of_first_instruction = list(dis.Bytecode(function))[0].starts_line or line_of_def_start + offset = line_of_first_instruction - line_of_def_start + return dedent("".join(lines[offset:])) def count_leading_whitespace(s: str) -> int: