Some Code Cleanup

Author: sarangbishal

examples/fibonacci.py

@@ -16,8 +16,9 @@ def fib(n, node_num):
     right = fib(n=n - 2, node_num=vs.node_count)
     return left + right
 
-# Call function
-print(fib(n=6, node_num=0))
+if __name__ == "__main__":
+    # Call function
+    print(fib(n=6, node_num=0))
 
-# Save recursion tree to a file
-vs.write_image("fibonacci.png")
+    # Save recursion tree to a file
+    vs.write_image("fibonacci.png")

visualiser/fibonacci.png

@@ -1,3 +1,5 @@
+# Author: Bishal Sarang
+
 import sys
 from functools import wraps
 from collections import OrderedDict
@@ -14,96 +16,104 @@ def __init__(self, ignore_args, show_argument_name=True, show_return_value=True)
         self.ignore_args = ignore_args
 
     @classmethod
-    def write_image(self, filename="out.png"):
-        self.graph.write_png(f"{filename}")
+    def write_image(cls, filename="out.png"):
+        try:
+            cls.graph.write_png(f"{filename}")
+            print(f"File {filename} successfully written")
+        except Exception:
+            print(f"Writing {filename} failed")
+
+    def extract_signature_label_arg_string(self, *args, **kwargs):
+        """
+        Returns function signature arguments as string and function label arguments as string.
+            label_argument string contains only the arguments that are not in ignore_args
+            signature_argument_string contains all the arguments available for the function
+        """
+        # If show_argument flag is True(default)
+        # Then argument_string is:
+        # a=1, b=31, c=0
+        signature_argument_string = ', '.join([repr(a) for a in args] +
+                                    [f"{key}={repr(value)}" for key, value in kwargs.items()])
+
+        label_argument_string = ', '.join([repr(a) for a in args] +
+                                                           [f"{key}={repr(value)}"
+                                                            for key, value in kwargs.items()
+                                                            if key not in self.ignore_args])
+
+        # If show_argument flag is False
+        # Then argument_string is:
+        # 1, 31, 0
+        if not self.show_argument_name:
+            signature_argument_string = ', '.join([repr(value) for value in args] +
+                                                 [f"{repr(value)}" for key, value in kwargs.items()])
+
+            label_argument_string = ', '.join([repr(a) for a in args] +
+                                              [f"{repr(value)}"
+                                               for key, value in kwargs.items()
+                                               if key not in self.ignore_args])
+
+        return signature_argument_string, label_argument_string
 
     def __call__(self, fn):
         @wraps(fn)
         def wrapper(*args, **kwargs):
             # Order all the keyword arguments
             kwargs = OrderedDict(sorted(kwargs.items()))
 
-            # If show_argument flag is True(default)
-            # Then argument_string is:
-            # a=1, b=31, c=0
-            argument_string = ', '.join(
-                [repr(a) for a in args] +
-                [f"{key}={repr(value)}" for key, value in kwargs.items()])
-
-            current_function_label_argument_string = ', '.join(
-                [repr(a) for a in args] +
-                [f"{key}={repr(value)}" for key, value in kwargs.items() if key not in self.ignore_args])
-
-            # If show_argument flag is False
-            # Then argument_string is:
-            # 1, 31, 0
-            if self.show_argument_name == False:
-                argument_string = ', '.join(
-                    [repr(value) for value in args] +
-                    [f"{repr(value)}" for key, value in kwargs.items()])
-
-                current_function_label_argument_string = ', '.join(
-                    [repr(a) for a in args] +
-                    [f"{repr(value)}" for key, value in kwargs.items() if key not in self.ignore_args])
+            """Details about current Function"""
+            # Get signature and label arguments strings for current function
+            current_function_argument_string, current_function_label_argument_string = self.extract_signature_label_arg_string(*args, **kwargs)
 
             # Details about current function
             current_function_name = fn.__name__
-            current_function_argument_string = argument_string
+
             # Current function signature looks as follows:
             # foo(1, 31, 0) or foo(a=1, b=31, c=0)
-            current_function_signature = f"{current_function_name}({current_function_argument_string})"
+            current_function_signature = f"{current_function_name}(" \
+                                         f"{current_function_argument_string})"
             current_function_label = f"{current_function_name}({current_function_label_argument_string})"
-
-            # Caller Function
-            caller_function_name = sys._getframe(1).f_code.co_name
-            # Extract the names of arguments only
-            caller_function_argument_names = sys._getframe(1).f_code.co_varnames[:fn.__code__.co_argcount]
-            caller_function_locals = sys._getframe(1).f_locals
-
+            """"""
+
+            """Details about caller function"""
+            caller_function_frame = sys._getframe(1)
+            # All the argument names in caller/parent function
+            caller_function_argument_names = caller_function_frame.f_code.co_varnames[
+                                             :fn.__code__.co_argcount]
+            caller_function_locals = caller_function_frame.f_locals
             # Sort all the locals of caller function
             caller_function_locals = OrderedDict(sorted(caller_function_locals.items()))
+            # Extract only those locals that are in arguments
+            caller_function_kwargs = {key: value for key, value in caller_function_locals.items() if key in caller_function_argument_names}
 
+            caller_function_argument_string, caller_function_label_argument_string = self.extract_signature_label_arg_string(**caller_function_kwargs)
 
-            # Extract only those locals which are in function signature
-            caller_function_argument_string = ', '.join(
-                [f"{key}={value}" for key, value in caller_function_locals.items() if (key in caller_function_argument_names)])
-            caller_function_label_argument_string = ', '.join(
-                [f"{key}={value}" for key, value in caller_function_locals.items() if
-                 (key in caller_function_argument_names and key not in self.ignore_args)])
-
-            if self.show_argument_name == False:
-                caller_function_argument_string = ', '.join(
-                    [f"{value}" for key, value in caller_function_locals.items() if
-                     (key in caller_function_argument_names)])
-                caller_function_label_argument_string = ', '.join(
-                    [f"{value}" for key, value in caller_function_locals.items() if
-                     (key in caller_function_argument_names and key not in self.ignore_args)])
+            # Caller Function
+            caller_function_name = caller_function_frame.f_code.co_name
 
+            # Extract the names of arguments only
             caller_func_signature = f"{caller_function_name}({caller_function_argument_string})"
             caller_func_label = f"{caller_function_name}({caller_function_label_argument_string})"
+            """"""
 
             if caller_function_name == '<module>':
                 print(f"Drawing for {current_function_signature}")
 
             result = fn(*args, **kwargs)
 
-            # #Child Node
-            child_label = current_function_label
+            # If show_return_value flag is set, display the result
             if self.show_return_value:
-                child_label += f" => {result}"
-            child_name = current_function_signature
-            v = pydot.Node(name=child_name, label=child_label)
-            self.graph.add_node(v)
+                current_function_label += f" => {result}"
 
-            # Parent Node
-            u = None
+            child_node = pydot.Node(name=current_function_signature, label=current_function_label)
+            self.graph.add_node(child_node)
 
+            # If the function is called by another function
             if caller_function_name != '<module>':
                 print(f"Called {current_function_label} by {caller_func_label}")
-                u = pydot.Node(name=caller_func_signature, label=caller_func_label)
-                self.graph.add_node(u)
-                edge = pydot.Edge(u, v)
-                self.graph.add_edge(edge)
 
+                parent_node = pydot.Node(name=caller_func_signature, label=caller_func_label)
+                self.graph.add_node(parent_node)
+                edge = pydot.Edge(parent_node, child_node)
+                self.graph.add_edge(edge)
             return result
         return wrapper