tfscript/src/tfscript/writing.py

import os
import sys
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=sys.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=sys.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})', file=sys.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=sys.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=sys.stderr)

    return list(fileNames.values())

def appendToActuals(targetDir, fileList):
    fileList = onlyFirsts(fileList)
    for currFile in fileList:
        execStr = f'exec {currFile.split(".")[0]}'
        realFilePath = targetDir + '/' + getRealName(currFile)

        realExists = exists(realFilePath)

        # creates if it doesn't exist, so must come after the exists() call
        cfgFile = open(realFilePath, 'a')

        if (realExists == False or not strInFile(execStr, cfgFile)):
            # since the file was created if it didn't exist,
            # and it was opened in append mode, this will
            # always place it at the end
            cfgFile.write( execStr.encode('utf8') )

        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"
    }
    if className in targetNames:
        className = targetNames[className]

    return className + '.cfg'

def strInFile(execStr, f):
    while True:
        line = f.readline()
        if line == "":
            # eof
            break
        # Remove indent and outdent, including trailing newline
        if execStr in line.strip():
            return True

    return False