From b1159d6c3eee51e846cdbd92f7ffc9ccdc7660c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?B=C3=A1lint=20Bokros?= <balint.bokros@avatao.com>
Date: Mon, 27 Nov 2017 18:51:20 +0100
Subject: [PATCH] Create helper methods for parsing source code

---
 lib/source_code.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)
 create mode 100644 lib/source_code.py

diff --git a/lib/source_code.py b/lib/source_code.py
new file mode 100644
index 0000000..9edef30
--- /dev/null
+++ b/lib/source_code.py
@@ -0,0 +1,26 @@
+import ast
+import inspect
+import re
+from io import StringIO
+
+
+def find_local_variable_value(func, local_variable_name):
+    func_src = inspect.getsource(func)
+    func_ast = ast.parse(func_src)
+    for node in ast.walk(func_ast):
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name):
+                    if target.id == local_variable_name:
+                        return node.value.s
+
+
+def get_source_code(func, strip_comments=False):
+    source = inspect.getsource(func)
+    if strip_comments:
+        # TODO: less fragile way to do this (tokenizer sadly inserts whitespaces all around)
+        comment_pattern = re.compile('^(\s.*)#.*$')
+        source = ''.join((
+            line for line in StringIO(source).readlines() if re.match(comment_pattern, line) is None
+        ))
+    return source