Fixed bug in parser that was preventing different argparse import styles from being recognized

added ast_inspector.py - utility for inspecting the properties of ast objects

 moved additional code formatting tasks into code_prep.py
 added code_prep_unittest.py

 Updated monkey_parser.py to use new parser fixes

gooey/code_prep.py

@@ -1 +1,64 @@
 __author__ = 'Chris'
+
+"""
+Preps the extracted Python code so that it can be evaled by the
+monkey_parser
+"""
+
+from itertools import *
+
+source = '''
+import sys
+import os
+import doctest
+import cProfile
+import pstats
+from argparse import ArgumentParser
+from argparse import RawDescriptionHelpFormatter
+from gooey import Gooey
+parser = ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
+parser.add_argument('filename', help='filename')
+parser.add_argument('-r', '--recursive', dest='recurse', action='store_true', help='recurse into subfolders [default: %(default)s]')
+parser.add_argument('-v', '--verbose', dest='verbose', action='count', help='set verbosity level [default: %(default)s]')
+parser.add_argument('-i', '--include', action='append', help='only include paths matching this regex pattern. Note: exclude is given preference over include. [default: %(default)s]', metavar='RE')
+parser.add_argument('-m', '--mycoolargument', help='mycoolargument')
+parser.add_argument('-e', '--exclude', dest='exclude', help='exclude paths matching this regex pattern. [default: %(default)s]', metavar='RE')
+parser.add_argument('-V', '--version', action='version')
+parser.add_argument('-T', '--tester', choices=['yes', 'no'])
+parser.add_argument(dest='paths', help='paths to folder(s) with source file(s) [default: %(default)s]', metavar='path', nargs='+')
+'''
+
+def take_imports(code):
+  return takewhile(lambda line: 'import' in line, code)
+
+def drop_imports(code):
+  return dropwhile(lambda line: 'import' in line, code)
+
+def split_line(line):
+  # splits an assignment statement into varname and command strings
+  # in: "parser = ArgumentParser(description='Example Argparse Program')"
+  # out: "parser", "= parser = ArgumentParser(description='Example Argparse Program"
+  # take/dropwhile used to avoid splitting on multiple '=' signs
+  not_equal_sign = lambda x: x != '='
+  varname = ''.join(takewhile(not_equal_sign, line)).strip()
+  command = ''.join(dropwhile(not_equal_sign, line))[2:]
+  return varname, command
+
+def update_parser_varname(new_varname, code):
+  # lines = source.split('\n')[1:]
+  lines = filter(lambda x: x != '', code)
+
+  argparse_code = dropwhile(lambda line: 'import' in line, lines)
+  old_argparser_varname, _ = split_line(argparse_code.next())
+
+  updated_code = [line.replace(old_argparser_varname, new_varname)
+                  for line in lines]
+  return updated_code
+
+if __name__ == '__main__':
+  pass
+
+
+
+
+

gooey/code_prep_unittest.py

@@ -0,0 +1,57 @@
+__author__ = 'Chris'
+
+import unittest
+import code_prep
+
+class TestCodePrep(unittest.TestCase):
+
+  def test_split_line(self):
+    line = "parser = ArgumentParser(description='Example Argparse Program')"
+    self.assertEqual("parser", code_prep.split_line(line)[0])
+    self.assertEqual("= ArgumentParser(description='Example Argparse Program')", code_prep.split_line(line)[1])
+
+  def test_update_parser_varname_assigns_new_name_to_parser_var(self):
+    line = ["parser = ArgumentParser(description='Example Argparse Program')"]
+    self.assertEqual(
+      "jarser = ArgumentParser(description='Example Argparse Program')",
+      code_prep.update_parser_varname('jarser', line)[0]
+    )
+
+  def test_update_parser_varname_assigns_new_name_to_parser_var__multiline(self):
+    lines = '''
+import argparse
+from argparse import ArgumentParser
+parser = ArgumentParser(description='Example Argparse Program')
+parser.parse_args()
+    '''.split('\n')
+
+    self.assertEqual(
+      "jarser = ArgumentParser(description='Example Argparse Program')",
+      code_prep.update_parser_varname('jarser', lines)[2]
+    )
+
+  def test_take_imports_drops_all_non_imports_statements(self):
+    lines = '''
+import argparse
+from argparse import ArgumentParser
+parser = ArgumentParser(description='Example Argparse Program')
+parser.parse_args()
+    '''.split('\n')[1:]
+
+    self.assertEqual(2, len(list(code_prep.take_imports(lines))))
+    self.assertEqual('import argparse', list(code_prep.take_imports(lines))[0])
+
+  def test_drop_imports_excludes_all_imports_statements(self):
+    lines = '''
+import argparse
+from argparse import ArgumentParser
+parser = ArgumentParser(description='Example Argparse Program')
+parser.parse_args()
+    '''.split('\n')[1:]
+
+    self.assertEqual(2, len(list(code_prep.take_imports(lines))))
+    self.assertEqual('parser.parse_args()', list(code_prep.drop_imports(lines))[1])
+
+if __name__ == "__main__":
+  #import sys;sys.argv = ['', 'Test.testName']
+  unittest.main()

gooey/dev_utils/ast_inspector.py

@@ -1 +1,52 @@
+from gooey import source_parser
+
 __author__ = 'Chris'
+
+"""
+Pretty Printing util for inspecting the various ast objects
+"""
+
+import ast
+from _ast import Assign, Call
+
+def pretty_print(node, indent):
+  d = node.__dict__
+  for k, v in d.iteritems():
+    if isinstance(v, list):
+      print '-' * indent, k, ": "
+      for i in v:
+        pretty_print(i, indent + 2)
+    elif 'ast' in str(type(v)):
+      pretty_print(v, indent + 2)
+    else:
+      print '-' * indent, k, ": ", v
+
+
+if __name__ == '__main__':
+  lines = '''
+def main():
+  x = 1
+  y = 2
+  foo, doo = ("poo", "poo")
+  smarser = argparse.ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
+  random_junk = 123412353454356
+  smarser.add_argument("filename", help="Name of the file you want to read")  # positional'
+  smarser.add_argument("outfile", help="Name of the file where you'll save the output")  # positional
+  bar = x + y
+  baz = random_junk * 5
+'''
+
+  lines2 = '''
+def main():
+  try:
+    foo, doo = ("poo", "poo")
+    smarser = argparse.ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
+    smarser.add_argument("filename", help="Name of the file you want to read")  # positional'
+    smarser.add_argument("outfile", help="Name of the file where you'll save the output")  # positional
+    smarser.parse_args()
+  except:
+    pass
+'''
+  nodes = ast.parse(open(r'C:\Users\Chris\Dropbox\pretty_gui\Gooey\gooey\mockapplications\mockapp_import_argparse.py').read())
+
+  pretty_print(nodes, 1)

gooey/mockapplications/example_argparse_souce_in_main.py

@@ -22,7 +22,7 @@ import os
 
 from argparse import ArgumentParser
 from argparse import RawDescriptionHelpFormatter
-from gooey.gooey_decorator import Gooey
+from gooey import Gooey
 
 __all__ = []
 __version__ = 0.1

gooey/mockapplications/mockapp_import_argparse.py

@@ -1 +1,51 @@
-__author__ = 'Chris'
+'''
+Created on Dec 21, 2013
+
+@author: Chris
+'''
+import sys
+import hashlib
+from time import time as _time
+from time import sleep as _sleep
+import argparse
+
+from gooey import Gooey
+
+
+@Gooey
+def main():
+  my_cool_parser = argparse.ArgumentParser(description="Mock application to test Gooey's functionality")
+  my_cool_parser.add_argument("filename", help="Name of the file you want to read")  # positional
+  my_cool_parser.add_argument("outfile", help="Name of the file where you'll save the output")  # positional
+  my_cool_parser.add_argument('-c', '--countdown', default=10, type=int, help='sets the time to count down from')
+  my_cool_parser.add_argument("-s", "--showtime", action="store_true", help="display the countdown timer")
+  my_cool_parser.add_argument("-d", "--delay", action="store_true", help="Delay execution for a bit")
+  my_cool_parser.add_argument('--verbose', '-v', action='count')
+  my_cool_parser.add_argument("-o", "--obfuscate", action="store_true", help="obfuscate the countdown timer!")
+  my_cool_parser.add_argument('-r', '--recursive', choices=['yes', 'no'], help='Recurse into subfolders')
+  my_cool_parser.add_argument("-w", "--writelog", default="No, NOT whatevs", help="write log to some file or something")
+  my_cool_parser.add_argument("-e", "--expandAll", action="store_true", help="expand all processes")
+
+  print 'inside of main(), my_cool_parser =', my_cool_parser
+  args = my_cool_parser.parse_args()
+
+  print sys.argv
+  print args.countdown
+  print args.showtime
+
+  start_time = _time()
+  print 'Counting down from %s' % args.countdown
+  while _time() - start_time < args.countdown:
+    if args.showtime:
+      print 'printing message at: %s' % _time()
+    else:
+      print 'printing message at: %s' % hashlib.md5(str(_time())).hexdigest()
+    _sleep(.5)
+  print 'Finished running the program. Byeeeeesss!'
+
+# 	raise ValueError("Something has gone wrong! AHHHHHHHHHHH")
+
+if __name__ == '__main__':
+  # 	sys.argv.extend('asdf -c 5 -s'.split())
+  # 	print sys.argv
+  main()

gooey/monkey_parser.py

@@ -9,6 +9,7 @@ import types
 from parser_exceptions import ArgumentError
 from argparse import ArgumentParser
 from argparse import RawDescriptionHelpFormatter
+import code_prep
 
 
 class MonkeyParser(object):
@@ -44,20 +45,22 @@ class MonkeyParser(object):
     and instantiates it in a local variable by evaling the rest of the lines.
     Each subsequent line updates the local variable in turn.
     '''
-    new_source_code = self._format_source_with_new_varname('clients_parser', source_code)
-    # variable of the same name as the one passed into the format_source method.
-    # Used to hold the eval'd statements
+    imports = filter(lambda x: 'gooey' not in x, code_prep.take_imports(source_code))
+    arg_code = code_prep.drop_imports(source_code)
+    updated_source_code = code_prep.update_parser_varname('clients_parser', arg_code)
 
-    first_line = new_source_code.pop(0)
-    clients_parser, assignment = self._split_line(first_line)
+    for _import in imports:
+      exec(_import)
 
+    first_line = updated_source_code.pop(0)
+    clients_parser, assignment = code_prep.split_line(first_line)
     clients_parser = eval(assignment)
 
-    for line in new_source_code:
+    for line in updated_source_code:
       eval(line)
     return clients_parser
 
-  def _format_source_with_new_varname(self, variable_name, source):
+  def _format_source_with_new_varname(self, new_variable_name, source):
     '''
     'injects' the client code with a known variable name so that it
     can be `eval`d and assigned to a variable in the local code.
@@ -74,11 +77,13 @@ class MonkeyParser(object):
     source_code = source[:]
 
     first_line = source_code[0]
-    parser_variable, statement = self._split_line(first_line)
-    parser_variable = parser_variable.strip()
+    client_parser_variable, statement = self._split_line(first_line)
+
+    client_parser_variable = client_parser_variable.strip()
 
     for index, line in enumerate(source_code):
-      source_code[index] = line.replace(parser_variable, variable_name)
+      source_code[index] = line.replace(client_parser_variable, new_variable_name)
+    source_code.append('{0}.parse_args()'.format(new_variable_name))
     return source_code
 
   def _split_line(self, line):

gooey/source_parser.py

@@ -9,6 +9,7 @@ client code.
 
 import os
 import ast
+import _ast
 from itertools import chain
 
 import codegen
@@ -33,114 +34,56 @@ def parse_source_file(file_name):
 
   Variables:
     * nodes 									Primary syntax tree object
-    * _mainfunc_block					main() method as found in the ast nodes
-    * _try_blocks							Try/except/finally blocks found in the main method
-    * main_block 							The code block in which the ArgumentParser statements are located
-    *	argparse_assign_obj   	The assignment of the ArgumentParser (line 1 in example code)
-    * parser_nodes						Nodes which have parser references (lines 1-4 in example code)
+    *	argparse_assignments   	The assignment of the ArgumentParser (line 1 in example code)
+    * add_arg_assignments     Calls to add_argument() (lines 2-3 in example code)
     * parser_var_name					The instance variable of the ArgumentParser (line 1 in example code)
     * ast_source							The curated collection of all parser related nodes in the client code
   """
 
   nodes = ast.parse(_openfile(file_name))
 
-  _mainfunc_block = find_main(nodes)
-  _try_blocks = find_try_blocks(_mainfunc_block)
+  module_imports = get_nodes_by_instance_type(nodes, _ast.Import)
+  specific_imports = get_nodes_by_instance_type(nodes, _ast.ImportFrom)
 
-  nodes_to_search = chain([_mainfunc_block], _try_blocks)
+  assignment_objs = get_nodes_by_instance_type(nodes, _ast.Assign)
+  call_objects = get_nodes_by_instance_type(nodes, _ast.Call)
 
-  main_block = find_block_containing_argparse(nodes_to_search)
-
-  argparse_assign_obj = find_assignment_objects(main_block)
-  parser_nodes = find_parser_nodes(main_block)
-  full_ast_source = chain(argparse_assign_obj, parser_nodes)
-  return full_ast_source
+  argparse_assignments = get_nodes_by_containing_attr(assignment_objs, 'ArgumentParser')
+  add_arg_assignments  = get_nodes_by_containing_attr(call_objects, 'add_argument')
+  # parse_args_assignment = get_nodes_by_containing_attr(call_objects, 'parse_args')
 
+  ast_argparse_source = chain(
+    module_imports,
+    specific_imports,
+    argparse_assignments,
+    add_arg_assignments
+    # parse_args_assignment
+  )
+  # for i in ast_argparse_source:
+  #   print i
+  return ast_argparse_source
 
 def _openfile(file_name):
   with open(file_name, 'rb') as f:
     return f.read()
 
+def get_nodes_by_instance_type(nodes, object_type):
+  return [node for node in walk_tree(nodes) if isinstance(node, object_type)]
 
-def find_main(nodes):
-  code_block = _find_block(nodes, ast.FunctionDef, lambda node: node.name == 'main')
-  if code_block != None:
-    return code_block
-  else:
-    raise ParserError('Could not find main function')
-
-
-def find_try_blocks(nodes):
-  return _find_blocks(nodes, [ast.TryExcept, ast.TryFinally], lambda x: x)
-
-
-def find_imports(nodes):
-  return _find_blocks(nodes, ast.ImportFrom, lambda x: x.module == 'argparse')
-
-
-def _find_block(nodes, types, predicate):
-  blocks = _find_blocks(nodes, types, predicate)
-  return blocks[0] if blocks else None
-
-
-def _find_blocks(nodes, types, predicate):
-  _types = types if isinstance(types, list) else [types]
-  return [node
-            for node in nodes.body
-            if any([isinstance(node, _type) for _type in _types])
-            and predicate(node)]
-
-
-def find_block_containing_argparse(search_locations):
-  # Browses a collection of Nodes for the one containing the Argparse instantiation
-  for location in search_locations:
-    if has_argparse_assignment(location):
-      return location
-  raise ParserError("Could not locate AugmentParser.")
-
-
-def has_argparse_assignment(block):
-  # Checks a given node for presence of an ArgumentParser instantiation
-  argparse_assignment = _find_statement(block, has_instantiator, 'ArgumentParser')
-  return is_found(argparse_assignment)
-
-
-def find_assignment_objects(ast_block):
-  return _find_statement(ast_block, has_instantiator, 'ArgumentParser')
+def get_nodes_by_containing_attr(nodes, attr):
+  return [node for node in nodes if attr in walk_tree(node)]
 
-def find_parser_nodes(ast_block):
-  return _find_statement(ast_block, has_assignment, 'add_argument')
-
-
-def is_found(stmnt):
-  return len(stmnt)
-
-
-def _find_statement(block, predicate, name):
-  return [node for node in block.body
-          if predicate(node, name)]
-
-
-def has_instantiator(x, name):
-  # Checks if the astObject is one with an instantiation of the ArgParse class
-  return has_attr(name, lambda _name: x.value.func.id == _name)
-
-
-def has_assignment(node, name):
-  # Checks if the astObject contains a function with a name of name
-  return has_attr(name, lambda _name: node.value.func.attr == _name)
-
-
-def has_attr(name, attr_predicate):
-  try:
-    return attr_predicate(name)
-  except AttributeError as e:
-    return False  #  Wrong type. Ignore.
-
-
-def get_assignment_name(node):
-  # return the variable name to which ArgumentParser is assigned
-  return node.targets[0].id
+def walk_tree(node):
+  yield node
+  d = node.__dict__
+  for key, value in d.iteritems():
+    if isinstance(value, list):
+      for val in value:
+        for _ in walk_tree(val): yield _
+    elif 'ast' in str(type(value)):
+      for _ in walk_tree(value): yield _
+    else:
+      yield value
 
 
 def convert_to_python(ast_source):
@@ -149,7 +92,6 @@ def convert_to_python(ast_source):
   """
   return map(codegen.to_source, ast_source)
 
-
 def extract_parser(modulepath):
   ast_source = parse_source_file(modulepath)
   if ast_source:
@@ -157,19 +99,19 @@ def extract_parser(modulepath):
     return MonkeyParser(python_code)
   return None
 
-
 if __name__ == '__main__':
   filepath = os.path.join(os.path.dirname(__file__),
                           'mockapplications',
                           'example_argparse_souce_in_main.py')
 
   nodes = ast.parse(_openfile(filepath))
-
+  #
   ast_source = parse_source_file(filepath)
-  python_code = convert_to_python(ast_source)
-  parser = MonkeyParser(python_code)
-  factory = ActionSorter(parser._actions)
-  print factory._positionals
+  python_code = convert_to_python(list(ast_source))
+  for i in python_code: print i
+  # parser = MonkeyParser(python_code)
+  # factory = ActionSorter(parser._actions)
+  # print factory._positionals
 
 
 

gooey/source_parser_unittest.py

@@ -108,14 +108,27 @@ else: pass
     result = source_parser.find_try_blocks(nodes)
     self.assertEqual(list(), result)
 
-  def test_find_argparse(selfs):
+  def test_find_argparse_located_object_when_imported_by_direct_name(self):
     example_source = '''
-parser = ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
-parser.add_argument("filename", help="filename")
-parser.add_argument("-r", "--recursive", dest="recurse", action="store_true", help="recurse into subfolders [default: %(default)s]")
-parser.add_argument("-v", "--verbose", dest="verbose", action="count", help="set verbosity level [default: %(default)s]")
+def main():
+  parser = ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
+    '''
+    nodes = ast.parse(example_source)
+    main_node = source_parser.find_main(nodes)
+    self.assertEqual('main', main_node.name)
+    containing_block = source_parser.find_block_containing_argparse([main_node])
+    self.assertTrue(containing_block is not None)
+
+  def test_find_argparse_located_object_when_access_through_module_dot_notation(self):
+    example_source = '''
+def main():
+  parser = argparse.ArgumentParser(description='Example Argparse Program', formatter_class=RawDescriptionHelpFormatter)
     '''
     nodes = ast.parse(example_source)
+    main_node = source_parser.find_main(nodes)
+    self.assertEqual('main', main_node.name)
+    containing_block = source_parser.find_block_containing_argparse([main_node])
+    self.assertTrue(containing_block is not None)
 
   def test_find_argparse_locates_assignment_stmnt_in_main(self):
     nodes = ast.parse(source_parser._openfile(self._module_with_argparse_in_main))
@@ -178,6 +191,28 @@ parser.add_argument("filename", help="filename")
     nodes = ast.parse(source)
     self.assertFalse(source_parser.has_instantiator(nodes.body[1], 'add_argument'))
 
+  def test_parser_identifies_import_module(self):
+    source = '''
+import os
+import itertools
+from os import path
+    '''
+    import _ast
+    nodes = ast.parse(source)
+    module_imports = source_parser.get_nodes_by_instance_type(nodes, _ast.Import)
+    self.assertEqual(2, len(module_imports))
+
+  def test_parser_identifies_import_from(self):
+    source = '''
+import os
+import itertools
+from os import path
+from gooey.gooey_decorator import Gooey
+    '''
+    import _ast
+    nodes = ast.parse(source)
+    from_imports = source_parser.get_nodes_by_instance_type(nodes, _ast.ImportFrom)
+    self.assertEqual(2, len(from_imports))