from ast import parse, NodeVisitor, Name
# For error messages (student-facing) only
_NAMES = {
'Add': '+',
'And': 'and',
'Assert': 'assert',
'Assign': '=',
'AnnAssign': '=',
'AugAssign': 'op=',
'BitAnd': '&',
'BitOr': '|',
'BitXor': '^',
'Break': 'break',
'Recursion': 'recursive call',
'ClassDef': 'class',
'Continue': 'continue',
'Del': 'del',
'Delete': 'delete',
'Dict': '{...}',
'DictComp': '{...}',
'Div': '/',
'Ellipsis': '...',
'Eq': '==',
'ExceptHandler': 'except',
'ExtSlice': '[::]',
'FloorDiv': '//',
'For': 'for',
'FunctionDef': 'def',
'Filter': 'filter',
'GeneratorExp': '(... for ...)',
'Global': 'global',
'Gt': '>',
'GtE': '>=',
'If': 'if',
'IfExp': '...if...else...',
'Import': 'import',
'ImportFrom': 'from ... import ...',
'In': 'in',
'Index': '...[...]',
'Invert': '~',
'Is': 'is',
'IsNot': 'is not ',
'LShift': '<<',
'Lambda': 'lambda',
'List': '[...]',
'ListComp': '[...for...]',
'Lt': '<',
'LtE': '<=',
'Mod': '%',
'Mult': '*',
'NamedExpr': ':=',
'Nonlocal': 'nonlocal',
'Not': 'not',
'NotEq': '!=',
'NotIn': 'not in',
'Or': 'or',
'Pass': 'pass',
'Pow': '**',
'RShift': '>>',
'Raise': 'raise',
'Return': 'return',
'Set': '{ ... } (set)',
'SetComp': '{ ... for ... } (set)',
'Slice': '[ : ]',
'Starred': '',
'Str': 'str',
'Sub': '-',
'Subscript': '[]',
'Try': 'try',
'Tuple': '(... , ... )',
'UAdd': '+',
'USub': '-',
'While': 'while',
'With': 'with',
'Yield': 'yield',
'YieldFrom': 'yield from',
}
def check(source_file, checked_funcs, disallow, source=None):
"""Checks that AST nodes whose type names are present in DISALLOW
(an object supporting 'in') are not present in the function(s) named
CHECKED_FUNCS in SOURCE. By default, SOURCE is the contents of the
file SOURCE_FILE. CHECKED_FUNCS is either a string (indicating a single
name) or an object of some other type that supports 'in'. CHECKED_FUNCS
may contain __main__ to indicate an entire module. Prints reports of
each prohibited node and returns True iff none are found.
See ast.__dir__() for AST type names. The special node name 'Recursion'
checks for overtly recursive calls (i.e., calls of the form NAME(...) where
NAME is an enclosing def."""
return ExclusionChecker(disallow).check(source_file, checked_funcs, source)
class ExclusionChecker(NodeVisitor):
"""An AST visitor that checks that certain constructs are excluded from
parts of a program. ExclusionChecker(EXC) checks that AST node types
whose names are in the sequence or set EXC are not present. Its check
method visits nodes in a given function of a source file checking that the
indicated node types are not used."""
def __init__(self, disallow=()):
"""DISALLOW is the initial default list of disallowed
node-type names."""
self._disallow = set(disallow)
self._checking = False
self._errs = 0
def generic_visit(self, node):
if self._checking and type(node).__name__ in self._disallow:
self._report(node)
super().generic_visit(node)
def visit_Module(self, node):
if "__main__" in self._checked_funcs:
self._checking = True
self._checked_name = self._source_file
super().generic_visit(node)
def visit_Call(self, node):
if 'Recursion' in self._disallow and \
type(node.func) is Name and \
node.func.id in self._func_nest:
self._report(node, "should not be recursive")
self.generic_visit(node)
def visit_FunctionDef(self, node):
self._func_nest.append(node.name)
if self._checking:
self.generic_visit(node)
elif node.name in self._checked_funcs:
self._checked_name = "Function " + node.name
checking0 = self._checking
self._checking = True
super().generic_visit(node)
self._checking = checking0
self._func_nest.pop()
def _report(self, node, msg=None):
node_name = _NAMES.get(type(node).__name__, type(node).__name__)
if msg is None:
msg = "should not contain '{}'".format(node_name)
print("{} {}".format(self._checked_name, msg))
self._errs += 1
def errors(self):
"""Returns the number of number of prohibited constructs found in
the last call to check."""
return self._errs
def check(self, source_file, checked_funcs, disallow=None, source=None):
"""Checks that AST nodes whose type names are present in DISALLOW
(an object supporting the contains test) are not present in
the function(s) named CHECKED_FUNCS in SOURCE. By default, SOURCE
is the contents of the file SOURCE_FILE. DISALLOW defaults to the
argument given to the constructor (and resets that value if it is
present). CHECKED_FUNCS is either a string (indicating a single
name) or an object of some other type that supports 'in'.
CHECKED_FUNCS may contain __main__ to indicate an entire module.
Prints reports of each prohibited node and returns True iff none
are found.
See ast.__dir__() for AST type names. The special node name
'Recursion' checks for overtly recursive calls (i.e., calls of the
form NAME(...) where NAME is an enclosing def."""
self._checking = False
self._source_file = source_file
self._func_nest = []
if type(checked_funcs) is str:
self._checked_funcs = { checked_funcs }
else:
self._checked_funcs = set(checked_funcs)
if disallow is not None:
self._disallow = set(disallow)
if source is None:
with open(source_file, 'r', errors='ignore') as inp:
source = inp.read()
p = parse(source, source_file)
self._errs = 0
self.visit(p)
return self._errs == 0
