tfscript/tfscript/writing.py

from sys import stderr
from os.path import exists
from tempfile import NamedTemporaryFile
from shutil import move as moveRobust

def writeOutput(data, className, args) -> dict:
    """
    Write `data' to various files as needed, returning a dict of
    the temporary file names and their target destination names,
    not including the target directory
    """
    namesDict = {} # return dict

    # Variables
    lineList = [ l.encode('utf8') for l in data.splitlines() ]
    fileNum = 1
    bytesWritten = 0

    # Constants
    maxFileSize = 2 ** 20 # 1MiB maximum cfg file size
    filesNeeded = 1 + int( len(data)/maxFileSize )
    if args.debug:
        print( f'DEBUG: need {filesNeeded} files for {className}', file=stderr)

    FilNedLen = len(str(filesNeeded))
    # extra 4 bytes is leeway
    reservedSpace = len(f'{className}_script_{filesNeeded}.cfg') + 4

    # Initialize variables
    outfile = NamedTemporaryFile(prefix=className, delete=False)
    # I know % formatting is old-school and pylint hates it,
    # but "%*d" is the easiest way to left-pad with zeros
    # without hardcoding a number.
    namesDict.update({ outfile.name: '%s_script_%0*d.cfg' % (className, FilNedLen, fileNum) })

    while (fileNum <= filesNeeded and len(lineList) > 0):
        line = lineList.pop(0) + '\n'.encode('utf8')
        lineLen = len(line) # nice

        if bytesWritten + reservedSpace + lineLen > maxFileSize:
            outfile.write( ('exec %s_script_%0*d' % (className, FilNedLen, fileNum+1)).encode('utf8') )
            bytesWritten += reservedSpace
            if args.debug:
                print( f'DEBUG: Wrote {bytesWritten} bytes to {className} ({fileNum}/{filesNeeded})', file=stderr)

            outfile.close()
            outfile = NamedTemporaryFile(prefix=className, delete=False)

            fileNum += 1
            namesDict.update({ outfile.name: '%s_script_%0*d.cfg' % (className, FilNedLen, fileNum) })
            bytesWritten = 0

        outfile.write(line)
        bytesWritten += lineLen

    outfile.close() # the most-recent tempfile will not have been closed
    if args.debug:
        print( f'DEBUG: Wrote {bytesWritten} bytes to {className} ({fileNum}/{filesNeeded})', end='\n\n', file=stderr)

    return namesDict

def replaceFiles(targetDir, fileNames, args):
    for tmpName, realName in fileNames.items():
        if args.dry_run:
            if args.debug:
                print( f'DEBUG: {tmpName} would be {targetDir}{realName}.cfg', file=stderr)
        else:
            # using shutil.move() because it can move files across disk drives on windows
            moveRobust( tmpName, f'{targetDir}{realName}' )
            if args.debug:
                print( f'DEBUG: Created {targetDir}{realName}', file=stderr)

    if args.debug:
        # Break up the debug messages
        print(end='\n')

    return list(fileNames.values())

def appendToActuals(targetDir, fileList, defaultsGiven, args):
    if defaultsGiven:
        classList = [
            "scout",
            "soldier",
            "pyro",
            "demo",
            "engi",
            "heavy",
            "medic",
            "sniper",
            "spy"
        ]
        for cclass in classList:
            addCallIfUncalled('exec default_script_1', targetDir, cclass, args)

    fileList = onlyFirsts(fileList)
    for currFile in fileList:
        execStr = f'exec {currFile.split(".")[0]}'
        addCallIfUncalled(execStr, targetDir, currFile, args)

def addCallIfUncalled(execStr, targetDir, fileName, args):
    realFilePath = targetDir + getRealName(fileName)

    realExists = exists(realFilePath)

    # creates if it doesn't exist, so must come after the exists() call
    cfgFile = open(realFilePath, 'a+')
    if not realExists:
        if args.debug:
            print( f"DEBUG: Created {realFilePath}" )
        cfgFile.write(execStr + '\n')

    elif not strInFile(execStr, cfgFile):
        cfgFile.write('\n' + execStr + '\n')

    cfgFile.close()

def onlyFirsts(fileList):
    for i, fileName in enumerate(fileList):
        noExtension = fileName.split('.')[0]
        number = int(noExtension.split('_')[2])
        if number != 1:
            fileList.pop(i)

    return fileList

def getRealName(fileName):
    className = fileName.split('_')[0]
    targetNames = {
        "demo": "demoman",
        "heavy": "heavyweapons",
        "engi": "engineer",
        "default": "autoexec"
    }
    if className in targetNames:
        className = targetNames[className]

    return className + '.cfg'

def strInFile(execStr, f):
    # Opened in append mode, so cursor is at the end.
    # Must reopen to put cursor at the start.
    with open(f.name, 'r') as dupfile:
        lineList = [ ' '.join(line.split()) for line in dupfile.readlines() ]
    for line in lineList:
        # Remove indent and outdent, including trailing newline
        if execStr == line:
            return True

    return False