<commit_before>from antlr4 import *
from AbstractSyntaxTree import *
from Visitor import *

class VisitorCodeGenerator(Visitor):

    def __init__(self, symbolTable, outFile="out.p"):
        self.symbolTable = symbolTable
        self.current = 0
        self.lvalue = []
        self.outFile = open(outFile, 'w')

        self.p_types = {
            "address" : "a",
            "int"     : "i",
            "float"   : "r",
            "char"    : "c",
            "numeric" : "N",
            "type"    : "T",
            "boolean" : "b"
        }

        self.initializers = {
            "int"     : 0,
            "float"   : 0.0,
            "char"    : 'c',
            "boolean" : 0
        }

        self.bin_arithm_op = {
            "+"  : "add",
            "-"  : "sub",
            "*"  : "mul",
            "/"  : "div",
            "%"  : "MODULO"
        }

        self.bin_comp_op = {
            "<" : "les",
            ">" : "grt",
            "<=" : "leq",
            ">=" : "geq",
            "==" : "equ",
            "!=" : "neq",
        }

        # - => neg
        # =
        # and -> and
        # or -> or
        # ! -> not

    def getLabel(self):
        self.current += 1
        return "l" + str(self.current)

    def __del__(self):
        self.outFile.close()

    def visitProgramNode(self, node):
        # self.outFile.write("code\n")
        # self.outFile.write("ujp  lmain\n")
        self.outFile.write("ujp main\n")
        self.visitChildren(node)

    def visitIncludeNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitFunctionDeclarationNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitFunctionDefinitionNode(self, node):
        self.symbolTable.openScope(isFunctionScope=True, name=node.identifier)
        node.plabel = "function_{0}".format(node.identifier)

        self.outFile.write("{0}:\n".format(node.plabel))
        self.outFile.write("sep 1000\n") # track actual value

        scope = self.symbolTable.currentScope
        for variable in scope.addressedVariables:
            # distinguish arguments from local variables: arguments already placed on stack by caller
            self.outFile.write("ldc {0} 0\n".format(self.p_types[variable.typeInfo.basetype]))

        for i in range(1, len(node.children)): # exclude parameters node
            node.children[i].accept(self)

        # set stack pointer back after wastefully using
        self.outFile.write("ssp {0}\n".format(self.symbolTable.currentScope.getAddressCounter() + 1))

        self.outFile.write("retp\n")

        self.symbolTable.closeScope()

    def visitMainFunctionNode(self, node):
        # self.outFile.write("\nlmain:\n")
        self.symbolTable.openScope(isFunctionScope=True, name=node.identifier)
        self.outFile.write("main:\n")
        self.visitChildren(node)
        self.outFile.write("hlt\n")
        self.symbolTable.closeScope()

    def visitParametersNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitParameterNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitArgumentsNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitStatementsNode(self, node):
        # self.outFile.write("code\n")
        openedScope = False
        if not isinstance(node.parent, ASTFunctionDefinitionNode):
            self.symbolTable.openScope()
            openedScope = True
        self.visitChildren(node)
        if openedScope:
            self.symbolTable.closeScope()
        # for ttype in self.symbolTable.getVariables(node):
        #     self.outFile.write("ldc {0} {1}".format(self.p_type[ttype], self.initializers[ttype]))

    def visitReturnNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitIfNode(self, node):
        elseLabel = self.getLabel()
        afterLabel = self.getLabel()

        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(elseLabel))        # if top == false, jump over the 'then' code
        node.children[1].accept(self)                            # 'then'

        if len(node.children) == 3:                              # optional else
            self.outFile.write("ujp {0}\n".format(afterLabel))   # jump over the 'else' code if coming from 'then'
            self.outFile.write("{0}:\n".format(elseLabel))
            node.children[2].accept(self)                        # else
            self.outFile.write("{0}:\n".format(afterLabel))
        else:
            self.outFile.write("{0}:\n".format(elseLabel))

    def visitElseNode(self, node):
        self.visitChildren(node)

    def visitWhileNode(self, node):
        conditionLabel = self.getLabel()
        afterLabel = self.getLabel()
        self.outFile.write("{0}:\n".format(conditionLabel))
        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(afterLabel))       # if top == false, jump over the loop code
        node.children[1].accept(self)                            # loop code
        self.outFile.write("ujp {0}\n".format(conditionLabel))   # jump back to the condition
        self.outFile.write("{0}:\n".format(afterLabel))

    def visitDoWhileNode(self, node): #TODO: test this
        conditionLabel = self.getLabel()
        afterLabel = self.getLabel()

        node.children[1].accept(self)                            # loop code

        self.outFile.write("{0}:\n".format(conditionLabel))
        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(afterLabel))       # if top == false, jump over the loop code
        node.children[1].accept(self)                            # loop code
        self.outFile.write("ujp {0}\n".format(conditionLabel))   # jump back to the condition
        self.outFile.write("{0}:\n".format(afterLabel))

    def visitVariableDeclarationNode(self, node):
        # nothing here, handled by declaratorInitializer
        self.visitChildren(node)

    def visitDeclaratorInitializerNode(self, node):
        hasInitializer = False
        for child in node.children:
            if isinstance(child, ASTInitializerListNode):
                hasInitializer = True

        self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address))

        if hasInitializer:
            self.visitChildren(node)
        else:
            self.outFile.write("ldc {0} {1}\n".format(self.p_types[node.getType().basetype], self.initializers[node.getType().basetype]))

        self.outFile.write("sto {0}\n".format(self.p_types[node.getType().basetype]))

    def visitInitializerListNode(self, node):
        # children will cause data to be loaded onto the stack
        self.visitChildren(node)

    def visitIntegerLiteralNode(self, node):
        self.outFile.write("ldc i " + str(node.value) + "\n")

    def visitFloatLiteralNode(self, node):
        self.outFile.write("ldc f " + str(node.value) + "\n")

    def visitCharacterLiteralNode(self, node):
        self.outFile.write("ldc c " + str(node.value) + "\n")

    def visitStringLiteralNode(self, node):
        self.outFile.write("code string literal\n")

    def visitVariableNode(self, node):
        if self.lvalue and self.lvalue.pop():
            self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address)) # TODO
        else:
            self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address)) # TODO
            self.outFile.write("ind {0}\n".format(self.p_types[node.getType().basetype])) # TODO
            # self.outFile.write("ind {0}\n".format(self.p_types[node.getType().basetype])) # TODO

        self.visitChildren(node)

    def visitFunctionCallNode(self, node):
        self.outFile.write("mst 0\n") # organizational block
        self.visitChildren(node) # evaluate arguments
        self.outFile.write("cup 0 {0}\n".format(node.definitionNode.plabel)) # call user procedure

    def visitTernaryConditionalOperatorNode(self, node):
        self.outFile.write("code ternary cond\n")
        self.visitChildren(node)

    def visitSimpleAssignmentOperatorNode(self, node):
        # children of a = b: [ASTVariableNode, ExpressionNode]

        self.lvalue.append(True)
        node.children[0].accept(self)
        self.outFile.write("dpl a\n")
        node.children[1].accept(self) # put expression result on stack as the "return value" of this operation
        self.outFile.write("sto {0}\n".format(self.p_types[node.children[0].getType().basetype]))
        self.outFile.write("ind {0}\n".format(self.p_types[node.children[0].getType().basetype]))

    def visitLogicOperatorNode(self, node):
        self.visitChildren(node)
        self.outFile.write(str(node.logicOperatorType) + "\n")

    def visitComparisonOperatorNode(self, node):
        self.visitChildren(node)
        self.outFile.write(self.bin_comp_op[str(node.comparisonType)] + " " + self.p_types[node.children[0].getType().basetype] + "\n")

    def visitUnaryArithmeticOperatorNode(self, node):
        op = str(node.arithmeticType)
        self.visitChildren(node)
        if op == "-":
            self.outFile.write("neg " + self.p_types[node.children[0].getType().basetype] +"\n")
        elif op == "+":
            pass
        elif op == "++":

            # self.outFile.write("ldc i 1\n")
            # self.outFile.write("add i\n")
            pass
        elif op == "--":
            self.outFile.write("UNARY ARITHMETIC" + op + "\n")

    def visitAddressOfoperatorNode(self, node):
        self.outFile.write("code address of op\n")
        self.visitChildren(node)

    def visitDereferenceNode(self, node):
        self.outFile.write("code dereference op\n")
        self.visitChildren(node)

    def visitLogicalNotOperatorNode(self, node):
        self.outFile.write("code logical not op\n")
        self.visitChildren(node)
        self.outFile.write("not " + self.p_types[node.children[0].getType().basetype] + "\n")

    def visitArraySubscriptNode(self, node):
        self.outFile.write("code array subscript op\n")
        self.visitChildren(node)

    def visitBinaryArithmeticNode(self, node):
        self.visitChildren(node)
        self.outFile.write(self.bin_arithm_op[str(node.arithmeticType)] + " " + self.p_types[node.children[0].getType().basetype] + "\n")
<commit_msg>More of the same (still not functional)<commit_after>from antlr4 import *
from AbstractSyntaxTree import *
from Visitor import *

class VisitorCodeGenerator(Visitor):

    def __init__(self, symbolTable, outFile="out.p"):
        self.symbolTable = symbolTable
        self.current = 0
        self.lvalue = []
        self.outFile = open(outFile, 'w')

        self.p_types = {
            "address" : "a",
            "int"     : "i",
            "float"   : "r",
            "char"    : "c",
            "numeric" : "N",
            "type"    : "T",
            "boolean" : "b"
        }

        self.initializers = {
            "int"     : 0,
            "float"   : 0.0,
            "char"    : 'c',
            "boolean" : 0
        }

        self.bin_arithm_op = {
            "+"  : "add",
            "-"  : "sub",
            "*"  : "mul",
            "/"  : "div",
            "%"  : "MODULO"
        }

        self.bin_comp_op = {
            "<" : "les",
            ">" : "grt",
            "<=" : "leq",
            ">=" : "geq",
            "==" : "equ",
            "!=" : "neq",
        }

        # - => neg
        # =
        # and -> and
        # or -> or
        # ! -> not

    def getLabel(self):
        self.current += 1
        return "l" + str(self.current)

    def __del__(self):
        self.outFile.close()

    def visitProgramNode(self, node):
        # self.outFile.write("code\n")
        # self.outFile.write("ujp  lmain\n")
        self.outFile.write("ujp main\n")
        self.visitChildren(node)

    def visitIncludeNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitFunctionDeclarationNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitFunctionDefinitionNode(self, node):
        self.symbolTable.openScope(isFunctionScope=True, name=node.identifier)
        node.plabel = "function_{0}".format(node.identifier)

        self.outFile.write("{0}:\n".format(node.plabel))
        self.outFile.write("sep 1000\n") # track actual value

        scope = self.symbolTable.currentScope
        for variable in scope.addressedVariables:
            # distinguish arguments from local variables: arguments already placed on stack by caller
            self.outFile.write("ldc {0} 0\n".format(self.p_types[variable.typeInfo.basetype]))

        for i in range(1, len(node.children)): # exclude parameters node
            node.children[i].accept(self)

        # set stack pointer back after wastefully using
        self.outFile.write("ssp {0}\n".format(self.symbolTable.currentScope.getAddressCounter() + 1))

        self.outFile.write("retp\n")

        self.symbolTable.closeScope()

    def visitMainFunctionNode(self, node):
        # self.outFile.write("\nlmain:\n")
        self.symbolTable.openScope(isFunctionScope=True, name=node.identifier)
        self.outFile.write("main:\n")
        for i in range(len(self.symbolTable.currentScope.addressedVariables)):
            self.outFile.write("ldc {0} 0\n".format(self.p_types[self.symbolTable.currentScope.addressedVariables[i].typeInfo.basetype]))
        self.visitChildren(node)
        self.outFile.write("hlt\n")
        self.symbolTable.closeScope()

    def visitParametersNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitParameterNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitArgumentsNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitStatementsNode(self, node):
        # self.outFile.write("code\n")
        openedScope = False
        if not isinstance(node.parent, ASTFunctionDefinitionNode):
            self.symbolTable.openScope()
            openedScope = True
        self.visitChildren(node)
        if openedScope:
            self.symbolTable.closeScope()
        # for ttype in self.symbolTable.getVariables(node):
        #     self.outFile.write("ldc {0} {1}".format(self.p_type[ttype], self.initializers[ttype]))

    def visitReturnNode(self, node):
        # self.outFile.write("code\n")
        self.visitChildren(node)

    def visitIfNode(self, node):
        elseLabel = self.getLabel()
        afterLabel = self.getLabel()

        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(elseLabel))        # if top == false, jump over the 'then' code
        node.children[1].accept(self)                            # 'then'

        if len(node.children) == 3:                              # optional else
            self.outFile.write("ujp {0}\n".format(afterLabel))   # jump over the 'else' code if coming from 'then'
            self.outFile.write("{0}:\n".format(elseLabel))
            node.children[2].accept(self)                        # else
            self.outFile.write("{0}:\n".format(afterLabel))
        else:
            self.outFile.write("{0}:\n".format(elseLabel))

    def visitElseNode(self, node):
        self.visitChildren(node)

    def visitWhileNode(self, node):
        conditionLabel = self.getLabel()
        afterLabel = self.getLabel()
        self.outFile.write("{0}:\n".format(conditionLabel))
        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(afterLabel))       # if top == false, jump over the loop code
        node.children[1].accept(self)                            # loop code
        self.outFile.write("ujp {0}\n".format(conditionLabel))   # jump back to the condition
        self.outFile.write("{0}:\n".format(afterLabel))

    def visitDoWhileNode(self, node): #TODO: test this
        conditionLabel = self.getLabel()
        afterLabel = self.getLabel()

        node.children[1].accept(self)                            # loop code

        self.outFile.write("{0}:\n".format(conditionLabel))
        node.children[0].accept(self)                            # condition
        self.outFile.write("fjp {0}\n".format(afterLabel))       # if top == false, jump over the loop code
        node.children[1].accept(self)                            # loop code
        self.outFile.write("ujp {0}\n".format(conditionLabel))   # jump back to the condition
        self.outFile.write("{0}:\n".format(afterLabel))

    def visitVariableDeclarationNode(self, node):
        # nothing here, handled by declaratorInitializer
        self.visitChildren(node)

    def visitDeclaratorInitializerNode(self, node):
        hasInitializer = False
        for child in node.children:
            if isinstance(child, ASTInitializerListNode):
                hasInitializer = True

        self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address))

        if hasInitializer:
            self.visitChildren(node)
        else:
            self.outFile.write("ldc {0} {1}\n".format(self.p_types[node.getType().basetype], self.initializers[node.getType().basetype]))

        self.outFile.write("sto {0}\n".format(self.p_types[node.getType().basetype]))

    def visitInitializerListNode(self, node):
        # children will cause data to be loaded onto the stack
        self.visitChildren(node)

    def visitIntegerLiteralNode(self, node):
        self.outFile.write("ldc i " + str(node.value) + "\n")

    def visitFloatLiteralNode(self, node):
        self.outFile.write("ldc f " + str(node.value) + "\n")

    def visitCharacterLiteralNode(self, node):
        self.outFile.write("ldc c " + str(node.value) + "\n")

    def visitStringLiteralNode(self, node):
        self.outFile.write("code string literal\n")

    def visitVariableNode(self, node):
        if self.lvalue and self.lvalue.pop():
            self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address)) # TODO
        else:
            self.outFile.write("ldc a {0}\n".format(node.symbolInfo.address)) # TODO
            self.outFile.write("ind {0}\n".format(self.p_types[node.getType().basetype])) # TODO
            # self.outFile.write("ind {0}\n".format(self.p_types[node.getType().basetype])) # TODO

        self.visitChildren(node)

    def visitFunctionCallNode(self, node):
        self.outFile.write("mst 0\n") # organizational block
        self.visitChildren(node) # evaluate arguments
        self.outFile.write("cup 0 {0}\n".format(node.definitionNode.plabel)) # call user procedure

    def visitTernaryConditionalOperatorNode(self, node):
        self.outFile.write("code ternary cond\n")
        self.visitChildren(node)

    def visitSimpleAssignmentOperatorNode(self, node):
        # children of a = b: [ASTVariableNode, ExpressionNode]

        self.lvalue.append(True)
        node.children[0].accept(self)
        self.outFile.write("dpl a\n")
        node.children[1].accept(self) # put expression result on stack as the "return value" of this operation
        self.outFile.write("sto {0}\n".format(self.p_types[node.children[0].getType().basetype]))
        self.outFile.write("ind {0}\n".format(self.p_types[node.children[0].getType().basetype]))

    def visitLogicOperatorNode(self, node):
        self.visitChildren(node)
        self.outFile.write(str(node.logicOperatorType) + "\n")

    def visitComparisonOperatorNode(self, node):
        self.visitChildren(node)
        self.outFile.write(self.bin_comp_op[str(node.comparisonType)] + " " + self.p_types[node.children[0].getType().basetype] + "\n")

    def visitUnaryArithmeticOperatorNode(self, node):
        op = str(node.arithmeticType)
        self.visitChildren(node)
        if op == "-":
            self.outFile.write("neg " + self.p_types[node.children[0].getType().basetype] +"\n")
        elif op == "+":
            pass
        elif op == "++":

            # self.outFile.write("ldc i 1\n")
            # self.outFile.write("add i\n")
            pass
        elif op == "--":
            self.outFile.write("UNARY ARITHMETIC" + op + "\n")

    def visitAddressOfoperatorNode(self, node):
        self.outFile.write("code address of op\n")
        self.visitChildren(node)

    def visitDereferenceNode(self, node):
        self.outFile.write("code dereference op\n")
        self.visitChildren(node)

    def visitLogicalNotOperatorNode(self, node):
        self.outFile.write("code logical not op\n")
        self.visitChildren(node)
        self.outFile.write("not " + self.p_types[node.children[0].getType().basetype] + "\n")

    def visitArraySubscriptNode(self, node):
        self.outFile.write("code array subscript op\n")
        self.visitChildren(node)

    def visitBinaryArithmeticNode(self, node):
        self.visitChildren(node)
        self.outFile.write(self.bin_arithm_op[str(node.arithmeticType)] + " " + self.p_types[node.children[0].getType().basetype] + "\n")
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__='Dan Tracy'
__email__='djt5019 at gmail dot com'

#This command line program will take a T.V. show as input and
#will return information about each episode, such as the title
#season and number.  I use this program to help clean up my
#TV show collection.  After the show has been found online it
#will be entered into a local database, provided sqlite3 is available,
#for faster lookup in the future.  You are able to filter the shows
#by season along with other options on the command line interface.
#You can also rename files according to a format that you choose as
#well as calculate the CRC32 of a file.  If you mistakenly rename
#files you have the option to revert the last renaming operation

import argparse
import os

import EpParser.src.Utils as Utils
import EpParser.src.Episode as Episode

from EpParser.src.Parser import EpParser as Parser
from EpParser.src.Cache import Cache
from EpParser.src.Logger import get_logger
from EpParser.src.Settings import Settings

def main():
    """ Our main function for our command line interface"""
    cmd = argparse.ArgumentParser(description="Renames your TV shows",
                                  prog='eplist', usage='%(prog)s [options] title')

    cmd.add_argument('title',
        help="The title of the show")

    cmd.add_argument('-d', '--display-header', action="store_true",
        help="Display the header at the top of the output")

    cmd.add_argument('-v', '--verbose', action="store_true",
        help="Be verbose, enable additional output")

    cmd.add_argument('-s', '--season', default="", type=str, metavar='N',
        help="The specific season range to search for. Ex: 1-3")

    cmd.add_argument('-e', '--episode', default="", type=str, metavar='N',
        help="The specific episode range to search for Ex: 15-30")

    cmd.add_argument('-f', '--format', dest="format", metavar='F',
        help="Rename the files in a directory with a custom format")

    cmd.add_argument('-g', '--gui-enabled', action="store_true",
        help="Use the gui rather than the command line")

    group = cmd.add_mutually_exclusive_group()

    group.add_argument('-r', '--rename', dest='pathname', metavar="PATH",
        help="Rename the files in the path provided")

    group.add_argument('-u', '--undo-rename', action='store_true',
        help="Undo the last rename operation")

    args = cmd.parse_args()

    if args.verbose:
        from logging import NOTSET
        for handle in get_logger().handlers:
            handle.setLevel(NOTSET)

    if args.gui_enabled:
        import EpParser.gui.gui as gui
        exit(gui.main())

    rename = args.pathname is not None

    if rename and not os.path.exists(args.pathname):
        exit("ERROR - Path provided does not exist")

    cache = Cache(Settings['episode_db'])
    episodeParser = Parser(args.title, cache)

    show = episodeParser.getShow()
    formatter = Episode.EpisodeFormatter(show, args.format)
    formatter.load_format_config()
    show.formatter = formatter

    if not show.episodeList:
        exit(1)

    # If the user specified a specific season we will filter our results
    # this also checks to make sure its a reasonable season number
    if args.season:
        seasonRange = list(parse_range(args.season))
        if seasonRange[-1] <= show.numSeasons:
            show.episodeList = [x for x in show.episodeList if x.season in seasonRange]

    if args.episode:
        episodeRange = list(parse_range(args.episode))

        if not args.season:
            show.episodeList = [x for x in show.episodeList if x.episodeCount in episodeRange]
        else:
            show.episodeList = [x for x in show.episodeList if x.episodeNumber in episodeRange]

    if args.undo_rename or rename:
        if rename:
            files = Utils.rename_files(args.pathname, show)
        else:
            files = Utils.load_last_renamed_files()

        if not files:
            print "Failed to find any files to rename"
            exit(1)

        print "PATH = {}".format(os.path.dirname(files[0][0]))
        for old, new in files:
            print (u"OLD: {0}".format(os.path.split(old)[1]))
            print (u"NEW: {0}".format(os.path.split(new)[1]))
            print

        errors = Utils.rename(files)
        if not errors:
            print "All files were successfully renamed"

        exit(0)

    if args.display_header or args.verbose:
        print "\nShow: {0}".format(args.title)
        print "Number of episodes: {0}".format(len(show.episodeList))
        print "Number of seasons: {0}".format(show.episodeList[-1].season)
        print "-" * 30

    print ""
    curr_season = show.episodeList[0].season
    for eps in show.episodeList:
        if curr_season != eps.season and args.display_header:
            print "\nSeason {0}".format(eps.season)
            print "----------"

        print show.formatter.display(eps)
        curr_season = eps.season

def parse_range(range):
    if '-' in range:
        high, low = range.split('-')
        high = int(high)
        low = int(low)
    else:
        high = low = int(range)

    low = max(low, 1)

    if low > high:
        low, high = high, low

    return xrange(low, high + 1)

if __name__ == '__main__':
    main()
<commit_msg>Went a little overboard with sed<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__='Dan Tracy'
__email__='djt5019 at gmail dot com'

#This command line program will take a T.V. show as input and
#will return information about each episode, such as the title
#season and number.  I use this program to help clean up my
#TV show collection.  After the show has been found online it
#will be entered into a local database, provided sqlite3 is available,
#for faster lookup in the future.  You are able to filter the shows
#by season along with other options on the command line interface.
#You can also rename files according to a format that you choose as
#well as calculate the CRC32 of a file.  If you mistakenly rename
#files you have the option to revert the last renaming operation

import argparse
import os

import EpParser.src.Utils as Utils
import EpParser.src.Episode as Episode

from EpParser.src.Parser import EpParser as Parser
from EpParser.src.Cache import Cache
from EpParser.src.Logger import get_logger
from EpParser.src.Settings import Settings

def main():
    """ Our main function for our command line interface"""
    cmd = argparse.ArgumentParser(description="Renames your TV shows",
                                  prog='eplist', usage='%(prog)s [options] title')

    cmd.add_argument('title',
        help="The title of the show")

    cmd.add_argument('-d', '--display-header', action="store_true",
        help="Display the header at the top of the output")

    cmd.add_argument('-v', '--verbose', action="store_true",
        help="Be verbose, enable additional output")

    cmd.add_argument('-s', '--season', default="", type=str, metavar='N',
        help="The specific season range to search for. Ex: 1-3")

    cmd.add_argument('-e', '--episode', default="", type=str, metavar='N',
        help="The specific episode range to search for Ex: 15-30")

    cmd.add_argument('-f', '--format', dest="format", metavar='F',
        help="Rename the files in a directory with a custom format")

    cmd.add_argument('-g', '--gui-enabled', action="store_true",
        help="Use the gui rather than the command line")

    group = cmd.add_mutually_exclusive_group()

    group.add_argument('-r', '--rename', dest='pathname', metavar="PATH",
        help="Rename the files in the path provided")

    group.add_argument('-u', '--undo-rename', action='store_true',
        help="Undo the last rename operation")

    args = cmd.parse_args()

    if args.verbose:
        from logging import NOTSET
        for handle in get_logger().handlers:
            handle.setLevel(NOTSET)

    if args.gui_enabled:
        import EpParser.gui.gui as gui
        exit(gui.main())

    rename = args.pathname is not None

    if rename and not os.path.exists(args.pathname):
        exit("ERROR - Path provided does not exist")

    cache = Cache(Settings['episode_db'])
    episodeParser = Parser(args.title, cache)

    show = episodeParser.getShow()
    formatter = Episode.EpisodeFormatter(show, args.format)
    formatter.load_format_config()
    show.formatter = formatter

    if not show.episodeList:
        exit(1)

    # If the user specified a specific season we will filter our results
    # this also checks to make sure its a reasonable season number
    if args.season:
        seasonRange = list(parse_range(args.season))
        if seasonRange[-1] <= show.numSeasons:
            show.episodeList = [x for x in show.episodeList if x.season in seasonRange]

    if args.episode:
        episodeRange = list(parse_range(args.episode))

        if not args.season:
            show.episodeList = [x for x in show.episodeList if x.episodeCount in episodeRange]
        else:
            show.episodeList = [x for x in show.episodeList if x.episodeNumber in episodeRange]

    if args.undo_rename or rename:
        if rename:
            files = Utils.rename_files(args.pathname, show)
        else:
            files = Utils.load_last_renamed_files()

        if not files:
            print "Failed to find any files to rename"
            exit(1)

        print "PATH = {}".format(os.path.dirname(files[0][0]))
        for old, new in files:
            print (u"OLD: {0}".format(os.path.split(old)[1]))
            print (u"NEW: {0}".format(os.path.split(new)[1]))
            print

        errors = Utils.rename(files)
        if not errors:
            print "All files were successfully renamed"

        exit(0)

    if args.display_header or args.verbose:
        print "\nShow: {0}".format(args.title)
        print "Number of episodes: {0}".format(len(show.episodeList))
        print "Number of seasons: {0}".format(show.episodeList[-1].season)
        print "-" * 30

    print ""
    curr_season = show.episodeList[0].season
    for eps in show.episodeList:
        if curr_season != eps.season and args.display_header:
            print "\nSeason {0}".format(eps.season)
            print "----------"

        print show.formatter.display(eps)
        curr_season = eps.season

def parse_range(range):
    if '-' in range:
        high, low = range.split('-')
        high = int(high)
        low = int(low)
    else:
        high = low = int(range)

    low = max(low, 1)

    if low > high:
        low, high = high, low

    return xrange(low, high + 1)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from enum import Enum

class IvlPortType(Enum):
    reg = 1
    wire = 2
    event = 3

class IvlDataDirection(Enum):
    input = 1
    output = 2

class IvlElabType(Enum):
    net_part_select = 1
    logic = 2
    posedge = 3
<commit_msg>Add docstrings to ivl_enums<commit_after>from enum import Enum

class IvlPortType(Enum):
    """
    Represents the type of an IvlPort.
    reg and wire are standard types declared in code.
    event is a type created by IVerilog when modules are linked to code events.
    """
    reg = 1
    wire = 2
    event = 3

class IvlDataDirection(Enum):
    """
    Represents the direction data flows to/from a part.
    Used by both IvlNet and IvlPort.
    """
    input = 1
    output = 2

class IvlElabType(Enum):
    """
    Represents the type of an IvlElab.
    """
    net_part_select = 1
    logic = 2
    posedge = 3
<|endoftext|>


<commit_before>import unittest

from vmmaster.core.virtual_machine.clone_factory import CloneFactory
from vmmaster.core.virtual_machine.clone import Clone
from vmmaster.core.config import setup_config, config
from vmmaster.core.platform import Platforms

class TestCloneFactory(unittest.TestCase):
    def setUp(self):
        setup_config('data/config.py')
        config.MAX_VM_COUNT = 5
        self.clone_factory = CloneFactory()
        platforms = Platforms()
        self.origin1 = platforms.platforms.values()[0]
        self.origin2 = platforms.platforms.values()[1]
        platforms.delete()

    def tearDown(self):
        self.clone_factory.delete()

    def test_clone_creation(self):
        clone = self.clone_factory.create_clone(self.origin1)
        self.assertIsInstance(clone, Clone)
        self.assertTrue(clone in self.clone_factory.clone_list.clones)

    def test_clone_deletion(self):
        clone = self.clone_factory.create_clone(self.origin1)
        clone.delete()
        self.assertFalse(clone in self.clone_factory.clone_list.clones)
        self.assertEqual(self.clone_factory.clone_list.total_count, 0)

    def test_creation_after_utilization(self):
        origin1_clones = [self.clone_factory.create_clone(self.origin1) for i in range(3)]
        origin2_clones = [self.clone_factory.create_clone(self.origin2) for i in range(2)]
        clones = origin1_clones + origin2_clones

        clone = clones.pop(0)
        clone.delete()
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))

        new_clone = self.clone_factory.create_clone(self.origin1)
        clones.append(new_clone)

        self.assertEqual(len(self.clone_factory.clone_list.clones), len(clones))
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))

    def test_clone_multiple_creation_and_multiple_utilization(self):
        origin1_clones = [self.clone_factory.create_clone(self.origin1) for i in range(3)]
        origin2_clones = [self.clone_factory.create_clone(self.origin2) for i in range(2)]
        clones = origin1_clones + origin2_clones

        from random import randint
        for i in range(3):
            number = randint(0, len(clones)-1)
            clone = clones.pop(number)
            clone.delete()

        for i in range(3):
            number = randint(0, 1)
            clone = self.clone_factory.create_clone(self.origin1 if number else self.origin2)
            clones.append(clone)

        self.assertEqual(len(self.clone_factory.clone_list.clones), len(clones))
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))<commit_msg>add clone factory delete test<commit_after>import unittest

from vmmaster.core.virtual_machine.clone_factory import CloneFactory
from vmmaster.core.virtual_machine.clone import Clone
from vmmaster.core.config import setup_config, config
from vmmaster.core.platform import Platforms

class TestCloneFactoryCreate(unittest.TestCase):
    def setUp(self):
        setup_config('data/config.py')
        config.MAX_VM_COUNT = 5
        self.clone_factory = CloneFactory()
        platforms = Platforms()
        self.origin1 = platforms.platforms.values()[0]
        self.origin2 = platforms.platforms.values()[1]
        platforms.delete()

    def tearDown(self):
        self.clone_factory.delete()

    def test_clone_creation(self):
        clone = self.clone_factory.create_clone(self.origin1)
        self.assertIsInstance(clone, Clone)
        self.assertTrue(clone in self.clone_factory.clone_list.clones)

    def test_clone_deletion(self):
        clone = self.clone_factory.create_clone(self.origin1)
        clone.delete()
        self.assertFalse(clone in self.clone_factory.clone_list.clones)
        self.assertEqual(self.clone_factory.clone_list.total_count, 0)

    def test_creation_after_utilization(self):
        origin1_clones = [self.clone_factory.create_clone(self.origin1) for i in range(3)]
        origin2_clones = [self.clone_factory.create_clone(self.origin2) for i in range(2)]
        clones = origin1_clones + origin2_clones

        clone = clones.pop(0)
        clone.delete()
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))

        new_clone = self.clone_factory.create_clone(self.origin1)
        clones.append(new_clone)

        self.assertEqual(len(self.clone_factory.clone_list.clones), len(clones))
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))

    def test_clone_multiple_creation_and_multiple_utilization(self):
        origin1_clones = [self.clone_factory.create_clone(self.origin1) for i in range(3)]
        origin2_clones = [self.clone_factory.create_clone(self.origin2) for i in range(2)]
        clones = origin1_clones + origin2_clones

        from random import randint
        for i in range(3):
            number = randint(0, len(clones)-1)
            clone = clones.pop(number)
            clone.delete()

        for i in range(3):
            number = randint(0, 1)
            clone = self.clone_factory.create_clone(self.origin1 if number else self.origin2)
            clones.append(clone)

        self.assertEqual(len(self.clone_factory.clone_list.clones), len(clones))
        self.assertEqual(sorted(self.clone_factory.clone_list.clones), sorted(clones))

class TestCloneFactoryDelete(unittest.TestCase):
    def setUp(self):
        setup_config('data/config.py')
        config.MAX_VM_COUNT = 5
        self.clone_factory = CloneFactory()
        platforms = Platforms()
        self.origin1 = platforms.platforms.values()[0]
        self.origin2 = platforms.platforms.values()[1]
        platforms.delete()

    def test_clone_factory_delete(self):
        clones = [self.clone_factory.create_clone(self.origin1) for i in range(3)]
        [clone.delete() for clone in clones]
        self.clone_factory.delete()
        self.assertEqual([], self.clone_factory.clone_list.clones)<|endoftext|>


<commit_before># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
from distutils import spawn
import logging
import os
import re
import select
import signal
import socket
import subprocess
import sys
import tenacity
import threading
import time

import fixtures
import jinja2
import six
try:
    import xattr
except ImportError:
    xattr = None

if six.PY3:
    fsdecode = os.fsdecode
else:
    def fsdecode(s):
        if isinstance(s, unicode):
            return s
        return s.decode(sys.getfilesystemencoding())

LOG = logging.getLogger(__name__)

class Driver(fixtures.Fixture):
    def __init__(self, env_prefix="PIFPAF", templatedir=".", debug=False,
                 tmp_rootdir=None):
        super(Driver, self).__init__()
        self.env_prefix = env_prefix
        self.env = {}
        self.debug = debug
        self.tmp_rootdir = tmp_rootdir

        templatedir = os.path.join('drivers', 'templates', templatedir)
        self.template_env = jinja2.Environment(
            loader=jinja2.PackageLoader('pifpaf', templatedir))

    def _setUp(self):
        self.tempdir = self.useFixture(fixtures.TempDir(self.tmp_rootdir)).path
        self.putenv("DATA", self.tempdir)

    @staticmethod
    def get_parser(parser):
        return parser

    def putenv(self, key, value, raw=False):
        if not raw:
            key = self.env_prefix + "_" + key
        self.env[key] = value
        return self.useFixture(fixtures.EnvironmentVariable(key, value))

    def _ensure_xattr_support(self):
        testfile = os.path.join(self.tempdir, "test")
        self._touch(testfile)
        xattr_supported = False
        if xattr is not None:
            try:
                x = xattr.xattr(testfile)
                x[b"user.test"] = b"test"
            except (OSError, IOError) as e:
                if e.errno != 95:
                    raise
            else:
                xattr_supported = True

        if not xattr_supported:
            raise RuntimeError("TMPDIR must support xattr for %s" %
                               self.__class__.__name__)

    def _kill(self, process):
        if not process.is_alive():
            return
        pgrp = os.getpgid(process.pid)
        process.terminate()
        try:

            self._wait(process)
        except tenacity.RetryError:
            LOG.warning("PID %d didn't terminate cleanly after 10 seconds, "
                        "sending SIGKILL to its process group", process.pid)
            # Cleanup remaining processes
            os.killpg(pgrp, signal.SIGKILL)
        process.wait()

    @staticmethod
    @tenacity.retry(wait=tenacity.wait_fixed(1),
                    stop=tenacity.stop_after_attempt(10),
                    retry=tenacity.retry_if_result(lambda ret: ret is None))
    def _wait(process):
        return process.poll()

    @staticmethod
    def find_executable(filename, extra_paths):
        paths = extra_paths + os.getenv('PATH', os.defpath).split(os.pathsep)
        for path in paths:
            loc = spawn.find_executable(filename, path)
            if loc is not None:
                return loc

    @staticmethod
    def find_config_file(filename):
        # NOTE(sileht): order matter, we first check into virtualenv
        # then global user installation, next system installation,
        # and to finish local user installation
        check_dirs = [sys.prefix + "/etc",
                      "/usr/local/etc",
                      "/etc",
                      os.path.expanduser("~/.local/etc")]
        for d in check_dirs:
            fullpath = os.path.join(d, filename)
            if os.path.exists(fullpath):
                return fullpath
        raise RuntimeError("Configuration file `%s' not found" % filename)

    def _read_in_bg(self, app, pid, fd):
        while True:
            data = fd.readline()
            if not data:
                break
            self._log_output(app, pid, data)

    @staticmethod
    def _log_output(appname, pid, data):
        data = fsdecode(data)
        LOG.debug("%s[%d] output: %s", appname, pid, data.rstrip())

    def _exec(self, command, stdout=False, ignore_failure=False,
              stdin=None, wait_for_line=None, wait_for_port=None,
              path=[], env=None,
              forbidden_line_after_start=None,
              allow_debug=True):
        LOG.debug("executing: %s" % command)

        app = command[0]

        debug = allow_debug and LOG.isEnabledFor(logging.DEBUG)

        if stdout or wait_for_line or debug:
            stdout_fd = subprocess.PIPE
        else:
            # TODO(jd) Need to close at some point
            stdout_fd = open(os.devnull, 'w')

        if stdin:
            stdin_fd = subprocess.PIPE
        else:
            # TODO(jd) Need to close at some point
            stdin_fd = open(os.devnull, 'r')

        if path or env:
            complete_env = dict(os.environ)
            if env:
                complete_env.update(env)
            if path:
                complete_env.update({
                    "PATH": ":".join(path) + ":" + os.getenv("PATH", ""),
                })
        else:
            complete_env = None

        try:
            c = subprocess.Popen(
                command,
                close_fds=True,
                stdin=stdin_fd,
                stdout=stdout_fd,
                stderr=subprocess.STDOUT,
                env=complete_env,
                preexec_fn=os.setsid,
            )
        except OSError as e:
            raise RuntimeError(
                "Unable to run command `%s': %s" % (b" ".join(command), e))

        self.addCleanup(self._kill, c)

        if stdin:
            LOG.debug("%s input: %s" % (app, stdin))
            c.stdin.write(stdin)
            c.stdin.close()

        if stdout or wait_for_line:
            lines = []
            while True:
                line = c.stdout.readline()
                self._log_output(app, c.pid, line)
                lines.append(line)
                if not line:
                    if wait_for_line:
                        raise RuntimeError(
                            "Program did not print: `%s'\nOutput: %s"
                            % (wait_for_line, b"".join(lines)))
                    break
                decoded_line = fsdecode(line)

                if wait_for_line and re.search(wait_for_line,
                                               decoded_line):
                    break
            stdout_str = b"".join(lines)
        else:
            stdout_str = None

        if (stdout or wait_for_line) and forbidden_line_after_start:
            timeout, forbidden_output = forbidden_line_after_start
            r, w, x = select.select([c.stdout.fileno()], [], [], timeout)
            if r:
                line = c.stdout.readline()
                self._log_output(app, c.pid, line)
                lines.append(line)
                if c.poll() is not None:
                    # Read the rest if the process is dead, this help debugging
                    while line:
                        line = c.stdout.readline()
                        self._log_output(app, c.pid, line)
                        lines.append(line)
                if line and re.search(forbidden_output, fsdecode(line)):
                    raise RuntimeError(
                        "Program print a forbidden line: `%s'\nOutput: %s"
                        % (forbidden_output, b"".join(lines)))

        if stdout or wait_for_line or debug:
            # Continue to read
            t = threading.Thread(target=self._read_in_bg,
                                 args=(app, c.pid, c.stdout,))
            t.setDaemon(True)
            t.start()

        if wait_for_port:
            for i in range(0, 10):
                with contextlib.closing(
                        socket.socket(socket.AF_INET,
                                      socket.SOCK_STREAM)) as sock:
                    if sock.connect_ex(('127.0.0.1', wait_for_port)) == 0:
                        break
                time.sleep(1)
            else:
                raise RuntimeError("Program did not open port %s" %
                                   wait_for_port)

        if not wait_for_line and not wait_for_port:
            status = c.wait()
            if not ignore_failure and status != 0:
                raise RuntimeError("Error while running command: %s" % command)

        return c, stdout_str

    def _touch(self, fname):
        open(fname, 'a').close()
        os.utime(fname, None)

    def template(self, resource, env, dest):
        template = self.template_env.get_template(resource)
        with open(dest, 'w') as f:
            f.write(template.render(**env))
<commit_msg>Fix process cleanup<commit_after># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
from distutils import spawn
import errno
import logging
import os
import re
import select
import signal
import socket
import subprocess
import sys
import tenacity
import threading
import time

import fixtures
import jinja2
import six
try:
    import xattr
except ImportError:
    xattr = None

if six.PY3:
    fsdecode = os.fsdecode
else:
    def fsdecode(s):
        if isinstance(s, unicode):
            return s
        return s.decode(sys.getfilesystemencoding())

LOG = logging.getLogger(__name__)

class Driver(fixtures.Fixture):
    def __init__(self, env_prefix="PIFPAF", templatedir=".", debug=False,
                 tmp_rootdir=None):
        super(Driver, self).__init__()
        self.env_prefix = env_prefix
        self.env = {}
        self.debug = debug
        self.tmp_rootdir = tmp_rootdir

        templatedir = os.path.join('drivers', 'templates', templatedir)
        self.template_env = jinja2.Environment(
            loader=jinja2.PackageLoader('pifpaf', templatedir))

    def _setUp(self):
        self.tempdir = self.useFixture(fixtures.TempDir(self.tmp_rootdir)).path
        self.putenv("DATA", self.tempdir)

    @staticmethod
    def get_parser(parser):
        return parser

    def putenv(self, key, value, raw=False):
        if not raw:
            key = self.env_prefix + "_" + key
        self.env[key] = value
        return self.useFixture(fixtures.EnvironmentVariable(key, value))

    def _ensure_xattr_support(self):
        testfile = os.path.join(self.tempdir, "test")
        self._touch(testfile)
        xattr_supported = False
        if xattr is not None:
            try:
                x = xattr.xattr(testfile)
                x[b"user.test"] = b"test"
            except (OSError, IOError) as e:
                if e.errno != 95:
                    raise
            else:
                xattr_supported = True

        if not xattr_supported:
            raise RuntimeError("TMPDIR must support xattr for %s" %
                               self.__class__.__name__)

    def _kill(self, process):
        process.terminate()
        try:

            self._wait(process)
        except tenacity.RetryError:
            LOG.warning("PID %d didn't terminate cleanly after 10 seconds, "
                        "sending SIGKILL to its process group", process.pid)
            # Cleanup remaining processes
            try:
                pgrp = os.getpgid(process.pid)
            except OSError as e:
                # ESRCH is returned if process just died in the meantime
                if e.errno != errno.ESRCH:
                    raise
            else:
                os.killpg(pgrp, signal.SIGKILL)
        process.wait()

    @staticmethod
    @tenacity.retry(wait=tenacity.wait_fixed(1),
                    stop=tenacity.stop_after_attempt(10),
                    retry=tenacity.retry_if_result(lambda ret: ret is None))
    def _wait(process):
        return process.poll()

    @staticmethod
    def find_executable(filename, extra_paths):
        paths = extra_paths + os.getenv('PATH', os.defpath).split(os.pathsep)
        for path in paths:
            loc = spawn.find_executable(filename, path)
            if loc is not None:
                return loc

    @staticmethod
    def find_config_file(filename):
        # NOTE(sileht): order matter, we first check into virtualenv
        # then global user installation, next system installation,
        # and to finish local user installation
        check_dirs = [sys.prefix + "/etc",
                      "/usr/local/etc",
                      "/etc",
                      os.path.expanduser("~/.local/etc")]
        for d in check_dirs:
            fullpath = os.path.join(d, filename)
            if os.path.exists(fullpath):
                return fullpath
        raise RuntimeError("Configuration file `%s' not found" % filename)

    def _read_in_bg(self, app, pid, fd):
        while True:
            data = fd.readline()
            if not data:
                break
            self._log_output(app, pid, data)

    @staticmethod
    def _log_output(appname, pid, data):
        data = fsdecode(data)
        LOG.debug("%s[%d] output: %s", appname, pid, data.rstrip())

    def _exec(self, command, stdout=False, ignore_failure=False,
              stdin=None, wait_for_line=None, wait_for_port=None,
              path=[], env=None,
              forbidden_line_after_start=None,
              allow_debug=True):
        LOG.debug("executing: %s" % command)

        app = command[0]

        debug = allow_debug and LOG.isEnabledFor(logging.DEBUG)

        if stdout or wait_for_line or debug:
            stdout_fd = subprocess.PIPE
        else:
            # TODO(jd) Need to close at some point
            stdout_fd = open(os.devnull, 'w')

        if stdin:
            stdin_fd = subprocess.PIPE
        else:
            # TODO(jd) Need to close at some point
            stdin_fd = open(os.devnull, 'r')

        if path or env:
            complete_env = dict(os.environ)
            if env:
                complete_env.update(env)
            if path:
                complete_env.update({
                    "PATH": ":".join(path) + ":" + os.getenv("PATH", ""),
                })
        else:
            complete_env = None

        try:
            c = subprocess.Popen(
                command,
                close_fds=True,
                stdin=stdin_fd,
                stdout=stdout_fd,
                stderr=subprocess.STDOUT,
                env=complete_env,
                preexec_fn=os.setsid,
            )
        except OSError as e:
            raise RuntimeError(
                "Unable to run command `%s': %s" % (b" ".join(command), e))

        self.addCleanup(self._kill, c)

        if stdin:
            LOG.debug("%s input: %s" % (app, stdin))
            c.stdin.write(stdin)
            c.stdin.close()

        if stdout or wait_for_line:
            lines = []
            while True:
                line = c.stdout.readline()
                self._log_output(app, c.pid, line)
                lines.append(line)
                if not line:
                    if wait_for_line:
                        raise RuntimeError(
                            "Program did not print: `%s'\nOutput: %s"
                            % (wait_for_line, b"".join(lines)))
                    break
                decoded_line = fsdecode(line)

                if wait_for_line and re.search(wait_for_line,
                                               decoded_line):
                    break
            stdout_str = b"".join(lines)
        else:
            stdout_str = None

        if (stdout or wait_for_line) and forbidden_line_after_start:
            timeout, forbidden_output = forbidden_line_after_start
            r, w, x = select.select([c.stdout.fileno()], [], [], timeout)
            if r:
                line = c.stdout.readline()
                self._log_output(app, c.pid, line)
                lines.append(line)
                if c.poll() is not None:
                    # Read the rest if the process is dead, this help debugging
                    while line:
                        line = c.stdout.readline()
                        self._log_output(app, c.pid, line)
                        lines.append(line)
                if line and re.search(forbidden_output, fsdecode(line)):
                    raise RuntimeError(
                        "Program print a forbidden line: `%s'\nOutput: %s"
                        % (forbidden_output, b"".join(lines)))

        if stdout or wait_for_line or debug:
            # Continue to read
            t = threading.Thread(target=self._read_in_bg,
                                 args=(app, c.pid, c.stdout,))
            t.setDaemon(True)
            t.start()

        if wait_for_port:
            for i in range(0, 10):
                with contextlib.closing(
                        socket.socket(socket.AF_INET,
                                      socket.SOCK_STREAM)) as sock:
                    if sock.connect_ex(('127.0.0.1', wait_for_port)) == 0:
                        break
                time.sleep(1)
            else:
                raise RuntimeError("Program did not open port %s" %
                                   wait_for_port)

        if not wait_for_line and not wait_for_port:
            status = c.wait()
            if not ignore_failure and status != 0:
                raise RuntimeError("Error while running command: %s" % command)

        return c, stdout_str

    def _touch(self, fname):
        open(fname, 'a').close()
        os.utime(fname, None)

    def template(self, resource, env, dest):
        template = self.template_env.get_template(resource)
        with open(dest, 'w') as f:
            f.write(template.render(**env))
<|endoftext|>


<commit_before>def crossCompartmentCorr(df, cyDict, cySets, studyStr, useFDR = False):
    """Plot of cytokine correlation for those that are common to both serum and ETT"""
    tmpCorr = zeros((len(cyDict[cySets[0]]),2))
    for i,(s,e) in enumerate(zip(cyDict[cySets[0]],cyDict[cySets[1]])):
        tmpCorr[i,0], tmpCorr[i,1] = partialcorr(df[s], df[e], method = 'spearman')
        #print s,e,tmpCorr[i,:]
        #print 'scatterfit(log(ndf["%s"]),log(ndf["%s"]),method="spearman")' % (s,e)
    sorti = argsort(tmpCorr[:,0])
    tmpCorr = tmpCorr[sorti,:]

    if useFDR:
        """Use q-value significance threshold"""
        sigInd, qvalues, _, _ = sm.stats.multipletests(tmpCorr[:,1], alpha = 0.2, method = 'fdr_bh')
    else:
        """Use p-value significance threshold"""
        sigInd = tmpCorr[:,1] < 0.05

    figure(111, figsize = (10,11.8))
    clf()
    barh(arange(tmpCorr.shape[0])[~sigInd], tmpCorr[~sigInd,0]**2, color = 'black', align='center')
    barh(arange(tmpCorr.shape[0])[sigInd], tmpCorr[sigInd,0]**2, color = 'red', align='center')
    yticks(range(tmpCorr.shape[0]), array([' '.join(cy.split(' ')[:-1]) for cy in cyDict[cySets[0]]])[sorti])
    grid(True, axis = 'x')
    xlabel('Correlation between\n%s and %s ($^*R^2$)' % (cySets[0],cySets[1]))
    ylim((-1,tmpCorr.shape[0]))
    xlim((0,1))
    tight_layout()
    figure(111).savefig(DATA_PATH + 'RandolphFlu/figures/%s_paired_corr.png' % studyStr)

def computeClusterOverlap(labelsA, labelsB):
    uA = unique(labelsA)
    uB = unique(labelsB)
    s = zeros((len(uA),len(uB)))
    for ai,a in enumerate(uA):
        for bi,b in enumerate(uB):
            indA = labelsA == a
            indB = labelsB == b
            s[ai,bi] = 2 * (indA & indB).sum()/(indA.sum() + indB.sum())
    sA = sorted(uA, key = lambda a: (labelsA == a).sum(), reverse = True)
    #sB = sorted(uB)
    sB = sorted(uB, key = lambda b: find(sA == uA[argmax(s[:,find(uB==b)])])[0])
    for ai,a in enumerate(sA):
        for bi,b in enumerate(sB):
            indA = labelsA == a
            indB = labelsB == b
            s[ai,bi] = 2 * (indA & indB).sum()/(indA.sum() + indB.sum())
    return sA, sB, s

def plotClusterOverlap(sA, sB, s):
    yA = linspace(10,0,len(sA))
    yB = linspace(10,0,len(sB))
    for ai, a in enumerate(sA):
        annotate(s = '%s' % a, xy = (0,yA[ai]),ha = 'center', va = 'center', size = 'x-large', zorder = 15)
        for bi, b in enumerate(sB):
            if ai == 0:
                annotate(s = '%s' % b, xy = (1,yB[bi]),ha = 'center', va = 'center', size = 'x-large', zorder = 15, color='white')
            plot([0,1],[yA[ai], yB[bi]], '-', lw = 20 * s[ai,bi], color='black', alpha = 0.7, zorder = 1)
    scatter(zeros(len(sA)), yA, s = 1000, color = 'red', zorder = 10)
    scatter(ones(len(sB)), yB, s = 1000, color = 'blue', zorder = 10)
    axis('off')

def alignClusters(predA,predB):
    """Returns a copy of predB with columns shuffled to maximize overlap with predA"""
    nCols = predA.shape[1]
    swap = {}
    for colA in range(nCols):
        match = argmax([(predA[:,colA] * predB[:,colB]).sum() for colB in range(nCols)])
        swap.update({match:colA})

    out = predB.copy()
    for colB,colA in swap.items():
        out[:,colA] = predB[:,colB]

    """In case the clusters aren't clearly 1:1 then try extra swaps until the optimum is found"""
    niters = 0
    while True:
        swaps = []
        curProd = (predA * out).sum()
        for ai,bi in itertools.product(range(nCols),range(nCols)):
            ind = arange(nCols)
            ind[ai] = bi
            ind[bi] = ai
            newProd = (predA * out[:,ind]).sum()
            if curProd < newProd:
                swaps.append((ai,bi,newProd))
        if len(swaps) == 0:
            break
        else:
            ai,bi,newProd = swaps[argmax([x[2] for x in swaps])]
            ind = arange(nCols)
            ind[ai] = bi
            ind[bi] = ai
            out = out[:,ind]
    return out<commit_msg>Hardened cluster comparison functions.<commit_after>import numpy as np
import matplotlib.pyplot as plt
import itertools
from sklearn import metrics
import pandas as pd
import statsmodels as sm
from corrplots import partialcorr

__all__ = ['compareClusters',
           'plotClusterOverlap',
           'alignClusters',
           'crossCompartmentCorr']

def compareClusters(labelsA, labelsB, method = 'ARI', alignFirst = True):
    if alignFirst:
        alignedB = alignClusters(labelsA, labelsB)
    else:
        alignedB = labelsB

    if method == 'ARI':
        s = metrics.adjusted_rand_score(labelsA.values, alignedB.values)
    elif method == 'AMI':
        s = metrics.adjusted_mutual_info_score(labelsA.values, alignedB.values)
    return s

def plotClusterOverlap(sA, sB, s):
    yA = np.linspace(10,0,len(sA))
    yB = np.linspace(10,0,len(sB))

    axh = plt.gca()
    axh.cla()
    annParams = dict(ha = 'center', va = 'center', size = 'x-large', zorder = 15)
    for ai, a in enumerate(sA):
        axh.annotate(s = '%s' % a, xy = (0,yA[ai]), color = 'black', **annParams)
        for bi, b in enumerate(sB):
            if ai == 0:
                annotate(s = '%s' % b, xy = (1,yB[bi]), color='white', **annParams)
            axh.plot([0,1], [yA[ai], yB[bi]], '-', lw = 20 * s[ai,bi], color='black', alpha = 0.7, zorder = 1)
    axh.scatter(np.zeros(len(sA)), yA, s = 1000, color = 'red', zorder = 10)
    axh.scatter(np.ones(len(sB)), yB, s = 1000, color = 'blue', zorder = 10)
    plt.axis('off')

def _alignClusterMats(matA, matB):
    """Returns a copy of matB with columns shuffled to maximize overlap with matA
    matX is a representation of cluster labels using a sparse\binary np.ndarray
    with labels along the columns"""
    nCols = matA.shape[1]
    swap = {}
    for colA in range(nCols):
        match = np.argmax([(matA[:,colA] * matB[:,colB]).sum() for colB in range(nCols)])
        swap.update({match:colA})

    out = matB.copy()
    for colB,colA in swap.items():
        out[:,colA] = matB[:,colB]

    """In case the clusters aren't clearly 1:1 then try extra swaps until the optimum is found"""
    niters = 0
    while True:
        swaps = []
        curProd = (matA * out).sum()
        for ai,bi in itertools.product(range(nCols),range(nCols)):
            ind = np.arange(nCols)
            ind[ai] = bi
            ind[bi] = ai
            newProd = (matA * out[:,ind]).sum()
            if curProd < newProd:
                swaps.append((ai,bi,newProd))
        if len(swaps) == 0:
            break
        else:
            ai,bi,newProd = swaps[np.argmax([x[2] for x in swaps])]
            ind = np.arange(nCols)
            ind[ai] = bi
            ind[bi] = ai
            out = out[:,ind]
    return out

def _alignSparseDf(dfA, dfB):
    out = _alignClusterMats(dfA.values, dfB.values)
    return pd.DataFrame(out, index = dfB.index, columns = dfB.columns)

def _labels2sparseDf(labels):
    labelCols = np.unique(labels)
    clusterMat = np.zeros((labels.shape[0], labelCols.shape[0]))
    for labi,lab in enumerate(labelCols):
        ind = (labels==lab).values
        clusterMat[ind,labi] = 1
    return pd.DataFrame(clusterMat, index = labels.index, columns = labelCols)

def _sparseDf2labels(sparseDf):
    labels = pd.Series(np.zeros(sparseDf.shape[0]), index = sparseDf.index)
    for clusterCol in sparseDf.columns:
        labels[sparseDf[clusterCol]] = clusterCol
    return labels

def alignClusters(labelsA, labelsB):
    """Returns a copy of labelsB with renamed labels shuffled to maximize overlap with matA"""
    sparseA = _labels2sparseDf(labelsA)
    sparseB = _labels2sparseDf(labelsB)
    outLabelsB = _sparseDf2labels(_alignSparseDf(sparseA, sparseB))
    return outLabelsB

def crossCompartmentCorr(dfA, dfB, method = 'pearson', useFDR = False, sigThreshold = 0.05):
    """Plot of cytokine correlation for those that are common to both A and B"""
    commonCy = [cy for cy in dfA.columns if cy in dfB.columns]
    tmpCorr = np.zeros((len(commonCy),2))
    for i,cy in enumerate(commonCy):
        tmpCorr[i,0], tmpCorr[i,1] = partialcorr(dfA[cy], dfB[cy], method = method)
    sorti = np.argsort(tmpCorr[:,0])
    tmpCorr = tmpCorr[sorti,:]

    if useFDR:
        """Use q-value significance threshold"""
        sigInd, qvalues, _, _ = sm.stats.multipletests(tmpCorr[:,1], alpha = sigThreshold, method = 'fdr_bh')
    else:
        """Use p-value significance threshold"""
        sigInd = tmpCorr[:,1] < sigThreshold

    axh = plt.gca()
    axh.cla()
    axh.barh(arange(tmpCorr.shape[0])[~sigInd], tmpCorr[~sigInd,0]**2, color = 'black', align='center')
    axh.barh(arange(tmpCorr.shape[0])[sigInd], tmpCorr[sigInd,0]**2, color = 'red', align='center')
    plt.yticks(range(tmpCorr.shape[0]), np.array([commonCy])[sorti])
    plt.grid(True, axis = 'x')
    plt.xlabel('Cross compartment correlation ($^*R^2$)')
    plt.ylim((-1,tmpCorr.shape[0]))
    plt.xlim((0,1))
    tight_layout()
<|endoftext|>


<commit_before>import time
from operator import itemgetter
from collections import OrderedDict

from IPython.display import ProgressBar

from jax.api_util import flatten_fun, wraps, flatten_fun_nokwargs
from jax.tree_util import tree_flatten, tree_unflatten, tree_map
import jax.core as jc
from jax.tree_util import build_tree
from jax import tree_util
import jax.interpreters.xla as xla
import jax.interpreters.ad as ad
from jax.ad_util import zeros_like_aval
import jax.interpreters.partial_eval as pe
import jax.linear_util as lu
from jax.util import safe_map, safe_zip, unzip2, partial, unzip3, cache
from jax import abstract_arrays
from jax import eval_shape

import fastar.util as util
from fastar.util import true_mask, false_mask

map = safe_map
zip = safe_zip

def get_aval(x):
    return abstract_arrays.raise_to_shaped(jc.get_aval(x))

class Env(OrderedDict): pass
def env_to_iterable(env):
    keys = tuple(env.keys())
    return tuple(map(env.get, keys)), keys
tree_util.register_pytree_node(Env, env_to_iterable,
                               lambda keys, xs: Env(zip(keys, xs)))

## Core
init_rules = {}
def init_ans(prim, *args, **params):
    if prim in init_rules:
        return init_rules[prim](*args, **params)
    else:
        args = [get_aval(arg) for arg, _ in args]
        abstract_out = prim.abstract_eval(*args, **params)
        if prim.multiple_results:
            return [parray(zeros_like_aval(o), false_mask(o))
                    for o in abstract_out]
        else:
            return parray(zeros_like_aval(abstract_out),
                          false_mask(abstract_out))

update_rules = {}
def get_update(p):
    try:
        return update_rules[p]
    except KeyError:
        raise NotImplementedError(
            "Fastar update rule for '{}' not implemented".format(p))

# Populate the cache
def firstpass(jaxpr, consts, freevar_vals, args):
    def read(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = env[v]
            assert isinstance(val, Parray)
            return val

    def write(v, val):
        assert isinstance(val, Parray)
        env[v] = val

    def write_subenvs(vs, env):
        subenvs[vs] = env

    env = Env()
    subenvs = Env()

    write(jc.unitvar, parray(jc.unit, true_mask(jc.unit)))
    map(write, jaxpr.constvars, consts)
    map(write, jaxpr.invars, args)
    map(write, jaxpr.freevars, freevar_vals)
    for eqn in jaxpr.eqns:
        in_vals = map(read, eqn.invars)
        if eqn.bound_subjaxprs:
            subjaxprs, sub_consts, sub_freevar_vals = unzip3([(
                subjaxpr,
                map(read, const_vars),
                map(read, bound_vars))
                for subjaxpr, const_vars, bound_vars
                in eqn.bound_subjaxprs])
            ans, subenvs_ = init_ans(eqn.primitive, eqn.params, subjaxprs,
                                     sub_consts, sub_freevar_vals, in_vals)
            ans, subenvs_ = get_update(eqn.primitive)(
                eqn.params, subjaxprs, sub_freevar_vals, in_vals, subenvs_)

            write_subenvs(tuple(eqn.outvars), subenvs_)
        else:
            ans = init_ans(eqn.primitive, *in_vals, **eqn.params)
            ans = get_update(eqn.primitive)(ans, *in_vals, **eqn.params)
        if eqn.primitive.multiple_results:
            map(write, eqn.outvars, ans)
        else:
            write(eqn.outvars[0], ans)
    return map(read, jaxpr.outvars), (env, subenvs)

def wrap(arr, mask):
    return parray(arr, mask.mask)

@lu.transformation_with_aux
def inited_fun(jaxpr, in_tree_def, masks, *args):
    args = [parray(a, m.mask) for a, m in zip(args, masks)]
    consts, freevar_vals, args = tree_unflatten(in_tree_def, args)
    out = yield (jaxpr, consts, freevar_vals, args), {}
    out = out[0], (out[1],)  # Tuple of envs, not singleton env
    out, out_mask, out_treedef = util.unmask_and_flatten(out)
    yield out, (out_mask, out_treedef)

def call_init_rule(primitive, params, jaxpr, consts, freevar_vals, in_vals):
    jaxpr, = jaxpr
    consts, = consts
    freevar_vals, = freevar_vals
    all_args, masks, in_treedef = util.unmask_and_flatten(
        (consts, freevar_vals, in_vals))
    masks = tuple(map(util.HashableMask, masks))
    fun = lu.wrap_init(firstpass)
    fun, out_mask_and_treedef = inited_fun(fun, jaxpr, in_treedef, masks)
    out = primitive.bind(fun, *all_args, **params)
    out_mask, out_treedef = out_mask_and_treedef()
    return tree_unflatten(out_treedef, map(parray, out, out_mask))
init_rules[xla.xla_call_p] = partial(call_init_rule, xla.xla_call_p)

def fastpass(jaxpr, freevar_vals, args, old_env):
    old_env, old_subenvs = old_env

    def read(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = env[v]
            assert isinstance(val, Parray)
            return val

    def read_old(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = old_env[v]
            assert isinstance(val, Parray)
            return val

    def read_old_subenvs(vs):
        return old_subenvs[vs]

    def write(v, val):
        assert isinstance(val, Parray)
        env[v] = val

    def write_subenvs(vs, env):
        subenvs[vs] = env

    def copy_old_to_new(v):
        write(v, read_old(v))

    env = Env()
    subenvs = Env()

    write(jc.unitvar, parray(jc.unit, true_mask(jc.unit)))
    map(copy_old_to_new, jaxpr.constvars)
    map(write, jaxpr.invars, args)
    map(write, jaxpr.freevars, freevar_vals)
    for eqn in jaxpr.eqns:
        in_vals = map(read, eqn.invars)
        old_outvals = map(read_old, eqn.outvars)
        if eqn.primitive.multiple_results:
            old_ans = old_outvals
        else:
            old_ans = old_outvals[0]
        in_vals = map(read, eqn.invars)
        if eqn.bound_subjaxprs:
            subjaxprs, sub_freevar_vals = unzip2([(
                subjaxpr,
                map(read, bound_vars))
                for subjaxpr, _, bound_vars in eqn.bound_subjaxprs])
            old_subenvs_ = read_old_subenvs(tuple(eqn.outvars))
            ans, subenvs_ = get_update(eqn.primitive)(
                eqn.params, subjaxprs, sub_freevar_vals, in_vals, old_subenvs_)
            write_subenvs(tuple(eqn.outvars), subenvs_)
        else:
            ans = get_update(eqn.primitive)(old_ans, *in_vals, **eqn.params)
        if eqn.primitive.multiple_results:
            map(write, eqn.outvars, ans)
        else:
            write(eqn.outvars[0], ans)
    return map(read, jaxpr.outvars), (env, subenvs)

@lu.transformation_with_aux
def updated_fun(jaxpr, in_treedef, masks, *args):
    args = [parray(a, m.mask) for a, m in zip(args, masks)]
    freevar_vals, args, env = tree_unflatten(in_treedef, args)
    out = yield (jaxpr, freevar_vals, args, env), {}
    out = out[0], (out[1],)  # Tuple of envs, not singleton env
    out, out_mask, out_treedef = util.unmask_and_flatten(out)
    yield out, (out_mask, out_treedef)

def call_update_rule(primitive, params, jaxpr, freevar_vals, in_vals, env):
    jaxpr, = jaxpr
    freevar_vals, = freevar_vals
    env, = env
    all_args, masks, in_treedef = util.unmask_and_flatten(
        (freevar_vals, in_vals, env))
    masks = tuple(map(util.HashableMask, masks))
    fun = lu.wrap_init(fastpass)
    fun, out_mask_and_treedef = updated_fun(fun, jaxpr, in_treedef, masks)
    out = primitive.bind(fun, *all_args, **params)
    out_mask, out_treedef = out_mask_and_treedef()
    return tree_unflatten(out_treedef, map(parray, out, out_mask))
update_rules[xla.xla_call_p] = partial(call_update_rule, xla.xla_call_p)

class Parray(tuple):
    """
    Array which has been partially evaluated, represented by a pair

        (arr, computed)

    where `computed` is a boolean array with True where arr has been computed,
    False where it hasn't yet. `arr` should be initialized to zeros so that
    np.all(~computed & arr == 0).
    """
    pass
parray = lambda arr, mask: Parray((arr, mask))

## High level API
def accelerate_part(fun):
    def first_fun(*args):
        ans, env = init_env(fun, args)
        def fast_fun(env, *args):
            ans, env = update_env(fun, args, env)
            return ans, partial(fast_fun, env)
        return ans, partial(fast_fun, env)
    return first_fun

@cache()
def fastar_jaxpr(fun, in_tree, in_avals):
    in_pvals = [pe.PartialVal((aval, jc.unit)) for aval in in_avals]
    fun_flat, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
    jaxpr, _, consts = pe.trace_to_jaxpr(fun_flat, in_pvals)
    return jaxpr, consts, out_tree()

def init_env(fun, args):
    args_flat, in_tree = tree_flatten(args)
    avals = tuple(get_aval(arg) for arg, _ in args_flat)
    jaxpr, consts, out_tree = fastar_jaxpr(fun, in_tree, avals)
    consts = [parray(const, true_mask(const)) for const in consts]
    ans, env = firstpass(jaxpr, consts, [], args_flat)
    return tree_unflatten(out_tree, ans), env

def update_env(fun, args, env):
    args_flat, in_tree = tree_flatten(args)
    avals = tuple(get_aval(arg) for arg, _ in args_flat)
    jaxpr, _, out_tree = fastar_jaxpr(fun, in_tree, avals)
    ans, env = fastpass(jaxpr, [], args_flat, env)
    return tree_unflatten(out_tree, ans), env

def accelerate(fixed_point_fun, max_iter=1000):
    def accelerated(x):
        x = tree_map(lambda x_leaf: parray(x_leaf, false_mask(x_leaf)), x)
        x, env = init_env(fixed_point_fun, [x])
        i = 1
        while not util.tree_mask_all(x) and i < max_iter:
            x, env = update_env(fixed_point_fun, [x], env)
            i = i + 1
        return tree_map(lambda x_leaf: x_leaf[0], x)
    return accelerated
<commit_msg>Rm progress-bar import<commit_after>import time
from operator import itemgetter
from collections import OrderedDict

from jax.api_util import flatten_fun, wraps, flatten_fun_nokwargs
from jax.tree_util import tree_flatten, tree_unflatten, tree_map
import jax.core as jc
from jax.tree_util import build_tree
from jax import tree_util
import jax.interpreters.xla as xla
import jax.interpreters.ad as ad
from jax.ad_util import zeros_like_aval
import jax.interpreters.partial_eval as pe
import jax.linear_util as lu
from jax.util import safe_map, safe_zip, unzip2, partial, unzip3, cache
from jax import abstract_arrays
from jax import eval_shape

import fastar.util as util
from fastar.util import true_mask, false_mask

map = safe_map
zip = safe_zip

def get_aval(x):
    return abstract_arrays.raise_to_shaped(jc.get_aval(x))

class Env(OrderedDict): pass
def env_to_iterable(env):
    keys = tuple(env.keys())
    return tuple(map(env.get, keys)), keys
tree_util.register_pytree_node(Env, env_to_iterable,
                               lambda keys, xs: Env(zip(keys, xs)))

## Core
init_rules = {}
def init_ans(prim, *args, **params):
    if prim in init_rules:
        return init_rules[prim](*args, **params)
    else:
        args = [get_aval(arg) for arg, _ in args]
        abstract_out = prim.abstract_eval(*args, **params)
        if prim.multiple_results:
            return [parray(zeros_like_aval(o), false_mask(o))
                    for o in abstract_out]
        else:
            return parray(zeros_like_aval(abstract_out),
                          false_mask(abstract_out))

update_rules = {}
def get_update(p):
    try:
        return update_rules[p]
    except KeyError:
        raise NotImplementedError(
            "Fastar update rule for '{}' not implemented".format(p))

# Populate the cache
def firstpass(jaxpr, consts, freevar_vals, args):
    def read(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = env[v]
            assert isinstance(val, Parray)
            return val

    def write(v, val):
        assert isinstance(val, Parray)
        env[v] = val

    def write_subenvs(vs, env):
        subenvs[vs] = env

    env = Env()
    subenvs = Env()

    write(jc.unitvar, parray(jc.unit, true_mask(jc.unit)))
    map(write, jaxpr.constvars, consts)
    map(write, jaxpr.invars, args)
    map(write, jaxpr.freevars, freevar_vals)
    for eqn in jaxpr.eqns:
        in_vals = map(read, eqn.invars)
        if eqn.bound_subjaxprs:
            subjaxprs, sub_consts, sub_freevar_vals = unzip3([(
                subjaxpr,
                map(read, const_vars),
                map(read, bound_vars))
                for subjaxpr, const_vars, bound_vars
                in eqn.bound_subjaxprs])
            ans, subenvs_ = init_ans(eqn.primitive, eqn.params, subjaxprs,
                                     sub_consts, sub_freevar_vals, in_vals)
            ans, subenvs_ = get_update(eqn.primitive)(
                eqn.params, subjaxprs, sub_freevar_vals, in_vals, subenvs_)

            write_subenvs(tuple(eqn.outvars), subenvs_)
        else:
            ans = init_ans(eqn.primitive, *in_vals, **eqn.params)
            ans = get_update(eqn.primitive)(ans, *in_vals, **eqn.params)
        if eqn.primitive.multiple_results:
            map(write, eqn.outvars, ans)
        else:
            write(eqn.outvars[0], ans)
    return map(read, jaxpr.outvars), (env, subenvs)

def wrap(arr, mask):
    return parray(arr, mask.mask)

@lu.transformation_with_aux
def inited_fun(jaxpr, in_tree_def, masks, *args):
    args = [parray(a, m.mask) for a, m in zip(args, masks)]
    consts, freevar_vals, args = tree_unflatten(in_tree_def, args)
    out = yield (jaxpr, consts, freevar_vals, args), {}
    out = out[0], (out[1],)  # Tuple of envs, not singleton env
    out, out_mask, out_treedef = util.unmask_and_flatten(out)
    yield out, (out_mask, out_treedef)

def call_init_rule(primitive, params, jaxpr, consts, freevar_vals, in_vals):
    jaxpr, = jaxpr
    consts, = consts
    freevar_vals, = freevar_vals
    all_args, masks, in_treedef = util.unmask_and_flatten(
        (consts, freevar_vals, in_vals))
    masks = tuple(map(util.HashableMask, masks))
    fun = lu.wrap_init(firstpass)
    fun, out_mask_and_treedef = inited_fun(fun, jaxpr, in_treedef, masks)
    out = primitive.bind(fun, *all_args, **params)
    out_mask, out_treedef = out_mask_and_treedef()
    return tree_unflatten(out_treedef, map(parray, out, out_mask))
init_rules[xla.xla_call_p] = partial(call_init_rule, xla.xla_call_p)

def fastpass(jaxpr, freevar_vals, args, old_env):
    old_env, old_subenvs = old_env

    def read(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = env[v]
            assert isinstance(val, Parray)
            return val

    def read_old(v):
        if type(v) is jc.Literal:
            return parray(v.val, true_mask(v.val))
        else:
            val = old_env[v]
            assert isinstance(val, Parray)
            return val

    def read_old_subenvs(vs):
        return old_subenvs[vs]

    def write(v, val):
        assert isinstance(val, Parray)
        env[v] = val

    def write_subenvs(vs, env):
        subenvs[vs] = env

    def copy_old_to_new(v):
        write(v, read_old(v))

    env = Env()
    subenvs = Env()

    write(jc.unitvar, parray(jc.unit, true_mask(jc.unit)))
    map(copy_old_to_new, jaxpr.constvars)
    map(write, jaxpr.invars, args)
    map(write, jaxpr.freevars, freevar_vals)
    for eqn in jaxpr.eqns:
        in_vals = map(read, eqn.invars)
        old_outvals = map(read_old, eqn.outvars)
        if eqn.primitive.multiple_results:
            old_ans = old_outvals
        else:
            old_ans = old_outvals[0]
        in_vals = map(read, eqn.invars)
        if eqn.bound_subjaxprs:
            subjaxprs, sub_freevar_vals = unzip2([(
                subjaxpr,
                map(read, bound_vars))
                for subjaxpr, _, bound_vars in eqn.bound_subjaxprs])
            old_subenvs_ = read_old_subenvs(tuple(eqn.outvars))
            ans, subenvs_ = get_update(eqn.primitive)(
                eqn.params, subjaxprs, sub_freevar_vals, in_vals, old_subenvs_)
            write_subenvs(tuple(eqn.outvars), subenvs_)
        else:
            ans = get_update(eqn.primitive)(old_ans, *in_vals, **eqn.params)
        if eqn.primitive.multiple_results:
            map(write, eqn.outvars, ans)
        else:
            write(eqn.outvars[0], ans)
    return map(read, jaxpr.outvars), (env, subenvs)

@lu.transformation_with_aux
def updated_fun(jaxpr, in_treedef, masks, *args):
    args = [parray(a, m.mask) for a, m in zip(args, masks)]
    freevar_vals, args, env = tree_unflatten(in_treedef, args)
    out = yield (jaxpr, freevar_vals, args, env), {}
    out = out[0], (out[1],)  # Tuple of envs, not singleton env
    out, out_mask, out_treedef = util.unmask_and_flatten(out)
    yield out, (out_mask, out_treedef)

def call_update_rule(primitive, params, jaxpr, freevar_vals, in_vals, env):
    jaxpr, = jaxpr
    freevar_vals, = freevar_vals
    env, = env
    all_args, masks, in_treedef = util.unmask_and_flatten(
        (freevar_vals, in_vals, env))
    masks = tuple(map(util.HashableMask, masks))
    fun = lu.wrap_init(fastpass)
    fun, out_mask_and_treedef = updated_fun(fun, jaxpr, in_treedef, masks)
    out = primitive.bind(fun, *all_args, **params)
    out_mask, out_treedef = out_mask_and_treedef()
    return tree_unflatten(out_treedef, map(parray, out, out_mask))
update_rules[xla.xla_call_p] = partial(call_update_rule, xla.xla_call_p)

class Parray(tuple):
    """
    Array which has been partially evaluated, represented by a pair

        (arr, computed)

    where `computed` is a boolean array with True where arr has been computed,
    False where it hasn't yet. `arr` should be initialized to zeros so that
    np.all(~computed & arr == 0).
    """
    pass
parray = lambda arr, mask: Parray((arr, mask))

## High level API
def accelerate_part(fun):
    def first_fun(*args):
        ans, env = init_env(fun, args)
        def fast_fun(env, *args):
            ans, env = update_env(fun, args, env)
            return ans, partial(fast_fun, env)
        return ans, partial(fast_fun, env)
    return first_fun

@cache()
def fastar_jaxpr(fun, in_tree, in_avals):
    in_pvals = [pe.PartialVal((aval, jc.unit)) for aval in in_avals]
    fun_flat, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
    jaxpr, _, consts = pe.trace_to_jaxpr(fun_flat, in_pvals)
    return jaxpr, consts, out_tree()

def init_env(fun, args):
    args_flat, in_tree = tree_flatten(args)
    avals = tuple(get_aval(arg) for arg, _ in args_flat)
    jaxpr, consts, out_tree = fastar_jaxpr(fun, in_tree, avals)
    consts = [parray(const, true_mask(const)) for const in consts]
    ans, env = firstpass(jaxpr, consts, [], args_flat)
    return tree_unflatten(out_tree, ans), env

def update_env(fun, args, env):
    args_flat, in_tree = tree_flatten(args)
    avals = tuple(get_aval(arg) for arg, _ in args_flat)
    jaxpr, _, out_tree = fastar_jaxpr(fun, in_tree, avals)
    ans, env = fastpass(jaxpr, [], args_flat, env)
    return tree_unflatten(out_tree, ans), env

def accelerate(fixed_point_fun, max_iter=1000):
    def accelerated(x):
        x = tree_map(lambda x_leaf: parray(x_leaf, false_mask(x_leaf)), x)
        x, env = init_env(fixed_point_fun, [x])
        i = 1
        while not util.tree_mask_all(x) and i < max_iter:
            x, env = update_env(fixed_point_fun, [x], env)
            i = i + 1
        return tree_map(lambda x_leaf: x_leaf[0], x)
    return accelerated
<|endoftext|>