improved maya ref parser

gistfile1.py

class MayaRefParser(object):
    '''
    This parses .ma file and get refrenced files info, store them in itself
    '''
    def __init__(self , maFile=None):
        self.header = ''
        self.refHeader = ''
        self.otherHeader = ''
        self.contents = ''
        self.contentsLineNum = 0
        self.maFile = maFile
        self.layoutIsLoaded = False
        if maFile and os.path.isfile(maFile):
            self.parse(maFile)
            # print self.refHeader
            self.parseRefHeader()

    def parse(self , maFile):
        '''
        Store header refHeader otherHeader and contentsLineNum after parsing
        '''
        with open(maFile) as inFile:

            inHeader = True
            inRefHeader = False
            inOtherHeader = False
            inContents = False
            for line in inFile:

                if inHeader and line.startswith('file -r'):
                    inHeader = False
                    inRefHeader = True
                elif inRefHeader and not line.startswith('file -r') and not line.startswith('\t'):
                    inRefHeader = False
                    inOtherHeader = True
                elif inOtherHeader and line.startswith('createNode'):
                    inOtherHeader = False
                    inContents = True
                if inHeader:
                    self.header += line
                elif inRefHeader:
                    self.refHeader += line
                elif inOtherHeader:
                    self.otherHeader += line
                elif inContents:
                    break

                self.contentsLineNum += 1

    def parseRefHeader(self):
        '''
        Parse refHeader and get referenced asset namespace / reference node / loaded status
        LAYOUT / SHAPELIB will be ignored
        '''
        import argparse
        import OrderedDict
        parser = argparse.ArgumentParser()
        parser.add_argument('-rfn')
        parser.add_argument('-rdi')
        parser.add_argument('-ns')
        parser.add_argument('-dr')
        # group = parser.add_mutually_exclusive_group()
        parser.add_argument('-r', action='store_true')
        parser.add_argument('file')
        parser.add_argument('-rpr')
        parser.add_argument('-shd',action='append')

        assetInfo = OrderedDict.OrderedDict()
        isLoaded = OrderedDict.OrderedDict()
        prefixRef = OrderedDict.OrderedDict()

        for line in self.refHeader.split(';\n')[:-1]:
            # print 'line -->',line
            args = parser.parse_args(line.split()[1:])
            if args.rdi:
                if args.ns:
                    # Namespace
                    nspace = args.ns.strip('"')
                # Asset Name
                aname = args.file.split("/")[-1].split(".")[0]
                # Reference Node
                refNode = args.rfn.strip('"')
                # Is reference loaded
                refIsLoaded = False if args.dr else True
                sharedRef = [ s.strip('"') for s in args.shd ] if args.shd else None

                if args.ns:
                    # Gather asset info
                    assetInfo[nspace] = {'asset':aname, 'refNode':refNode, 'refIsLoaded':refIsLoaded,'sharedRef':sharedRef}

                    if refIsLoaded:
                        isLoaded[nspace] = {'asset':aname, 'refNode':refNode, 'refIsLoaded':refIsLoaded,'sharedRef':sharedRef}

                # when prefix is used to create ref, there is no -ns
                else:

                    prefix = args.rpr.strip('"')
                    # Namespace
                    prefixRef[prefix] = {'asset':aname, 'refNode':refNode, 'refIsLoaded':refIsLoaded,'sharedRef':sharedRef}

            elif args.rdi and args.ns == '"LAYOUT"':
                self.layoutIsLoaded = False if args.dr else True

        self.assetInfo = assetInfo
        self.prefixRef = prefixRef
        self.isLoaded = isLoaded

    def loadRefAndWriteNewMayaFile(self, loadPackages=None, newFile=None):
        '''replace -dr 1 in .ma files first few lines , so ref will be loaded when it next time gets opened'''
        with open(newFile, 'w') as new_file:

            new_file.write(self.header)

            for line in self.refHeader.split(';\n')[:-1]:
                # rewrite ref header in single line , may be dangerous but easier to write regex
                newLine = line.replace('\n','').replace('\t','')

                # print 'line -->',line
                if ' -ns ' in line:
                    ns = re.search('(?<=-ns \").+(?=\" -)',newLine ).group()
                    if ns in loadPackages:
                        # print "match line : ", line
                        newLine = re.sub('(?<= )-dr 1 (?=-rfn)', '', newLine)
                        # print "new line   :",newLine
                    # else:
                        # print 'skip...','found no', ns
                new_file.write(newLine + ';\n')

            new_file.write(self.otherHeader)

            with open(self.maFile, 'r') as old_file:
                for _ in xrange(self.contentsLineNum):
                    old_file.next()
                for old_line in old_file:
                    # pass
                    new_file.write(old_line)