def makeNode(className, lvl, context): """ Given a string className, a tree-depth level, and a context stack, creates and returns a full AST node generated using the PCFG. All fields and children of this AST node are recursively populated as necessary. Args: className: string class name of the node that should be created, e.g. "Module" or "FunctionDef" lvl: current depth in the AST. The initial call to this function should have lvl=0. This is mainly used for help printing out logs of the tree creation, but can certainly be used in the future to do smarter generation in a context-sensitive way. context: a list, representing a stack, of ancestors in this AST call-chain. Each element of the stack is of the form (className, fieldName), e.g. ("FunctionDef", "args") meaning that this current node is a child of FunctionDef under its "args" field. """ log("makeNode({})".format(className),lvl) # copy object since `objects` just stores representative # object types node = copy.copy(objects[className]) # Get possible rules and draw a random rule from those given # the probabilities of each rule in the PCFG possible_rules = heads[className] probabilities = [pcfg[rule] for rule in possible_rules] sRule = util.random_draw(possible_rules, probabilities) (head, constituent) = util.sRule2tRule(sRule) log("rule: {}".format(sRule),lvl) if constituent != "<NULL>": # For each RHS element in the rule, populate the element # appropriately and set the attribute. for attrPair in constituent: newContext = context + [(className, attrPair[0])] populatedField = populateField(field=attrPair[1], lvl=lvl+1, context=newContext) setattr(node, attrPair[0], populatedField) return node
def make_primitive(primitive_className, context): """ Given the class name of a primitive, as well as the context stack, returns a randomly chosen primitive of the correct type following certain restrictions if applicable. Args: primitive_className: the classname of the primitive. Should be one of the seven class names given in util.PRIMITIVE_CLASSNAMES context: context stack """ # Recall primitives dict is of the form: # primitives[pcn]["normal"|"special"]["vals"|"freqs"] style = "special" if is_namelike(primitive_className, context) else "normal" # Randomly draw one primitive vals = primitives[primitive_className][style]["vals"] frequencies = primitives[primitive_className][style]["freqs"] return util.random_draw(vals, frequencies)