extract.py

"""Contains basic functions for updating name mangling extracted from dumpbin's output.
This is a simple tool that operates on text level without any syntax validation.

When encounters overloads for symbol will write all variants on separate lines,
then user have to pick correct line and remove the rest manually.
"""

from string import Template
from collections import namedtuple
import argparse
import re
import os
import sys


VARIANTS_TEMPLATE = Template('version(Windows) pragma(mangle, "$name") //$signature \n')
MANGLING_TEMPLATE = Template(#'version(Windows) pragma(mangle, "$mangle") //$signature \n'
                '  $retType $fnName$argList; else\n'
                '$retType $fnName$argList;\n\n')

MARKER_ATTR = '@pyExtract'

# Matches functions marked with attribute
REPLACESIG_RE = re.compile(r'@pyExtract.*\s+'
                           r'(?P<retType>\w+)\s+'
                           r'(?P<name>\w+)\s*'
                           r'(?P<args>\([\w"%.()*\s,/\[\]=-]+);')

# Match groups for 'mangled name (function signature)'
MANGLING_RE = re.compile(r'(?P<decorated>[\?\w@]+)'
                         r' \((?P<signature>[\(\)\w\s_:*&,]+)\)')


DecoratedSymbol = namedtuple('DecoratedSymbol', ['name', 'signature'])

MangleData = namedtuple('MangleData', ['retType', 'name', 'args', 'variants'])


generatedFileWarning = '// NOTE:\n' \
                '//   This file is automatically generated.\n' \
                '//   Any manual changes will be lost on next update.\n' \
                '\n'

# Contains line information for replacement after extracting @pyExtract markers
replaces = dict()

# Text to add before output content
prepend = ''

# Stores info from exports file about signatures with const-pointers 
cpexport = dict()
cpexportNames = list()


def report(message):
    """Reports progress to console if no 'quiet' attribute is found in this module"""
    if hasattr(sys.modules[__name__], 'quiet'):
        return
    print(message)

def append_suffix(dfile, suf='_gen'):
    """Appends suffix before .d extensions"""
    return dfile[0:dfile.rindex('.d')] + suf + '.d'


def process_mangling(source, exports=None, output=None):
    """Single entry point for extracting and writing generated mangling file""" 

    report('building exports lookup information')
    _extract_preload(exports)
    find_const_pointers(exports)

    report(f'scanning {source}')
    lines, linesString = scan_file(source)

    spanLines = map_lines(linesString, replaces.keys())
    lineSpans = {spanLines[k]:k for k in spanLines}

    if output is None:
        output = append_suffix(source)

    try:
        global prepend
        if os.path.isfile(prepend):
            with open(prepend) as pf:
                prepend = pf.readlines()
    except FileNotFoundError:
        pass

    report(f'writing generated data to {output}')
    with open(output, 'w') as outf:
        outf.write(generatedFileWarning)
        outf.write(prepend)
        for srcNum, srcLn in enumerate(lines):
            if srcNum in lineSpans.keys():   # Line where marker attribute found
                replData     =  replaces[lineSpans[srcNum]]
                numVariants  =  len(replData.variants)
                funcParts = dict(
                                fnName=replData.name,
                                retType=replData.retType,
                                argList=replData.args)
                # Write variants stacked on top of each other
                for i in range(numVariants): 
                    sigvariant = VARIANTS_TEMPLATE.substitute(replData.variants[i]._asdict())
                    outf.write(sigvariant)
                text = MANGLING_TEMPLATE.substitute(funcParts)
                outf.write(text)


def scan_file(fname):
    """Builds a list of @pyReplace's from file"""

    with open(fname) as f:
        lines = f.readlines()
        linesString = ''.join(lines)
        for it in REPLACESIG_RE.finditer(linesString):
            fnName   = it.group('name')
            start, _ = it.span()
            for m in extract_mangling(fnName):
                entry = replaces.get(start, MangleData(it.group('retType'), fnName, it.group('args'), list()))
                entry.variants.append(m)
                replaces[start] = entry
    return lines, linesString


def map_lines(text, spans):
    """Builds dictionary of spans to lines numbers"""

    lineno = dict()
    previous = 0
    lastline = 0
    for kl in spans:
        start = kl
        curline = text.count('\n', previous, start) + lastline
        lineno[start] = curline
        lastline = curline
        previous = start
    return lineno


def find_const_pointers(exportsFile):
    """Searches the exports file for signatures containing const-pointer-to-mutable-type"""

    CONSTPTR_RE = re.compile(r'\w* \* const') # "type * const"
    with open(exportsFile) as f:
        for lnum, line in enumerate(f):
            match = CONSTPTR_RE.search(line)
            if match is not None:
                cpexport[lnum] = line
        for _,ln in enumerate(cpexport):
            match = MANGLING_RE.search(cpexport[ln])
            if match is not None:
                sig    =  match.group('signature')
                nameStart  =  sig.rfind('::')
                if nameStart < 0:
                    continue
                nameStart  =  nameStart + 2
                nameEnd    =  sig.index('(')
                cpexportNames.append(sig[nameStart:nameEnd])


def _extract_preload(exportsFile):
    """Preloads file content for extract_mangling function"""
    if not hasattr(extract_mangling, 'exports'):
        with open(exportsFile) as f:
            extract_mangling.exports = f.readlines() 

def extract_mangling(sym):
    """Extracts mangled name(s) for specified symbol (all overloads)"""

    # Walks over ocurrences of sym in the file
    for i in (mangled for mangled in extract_mangling.exports if mangled.find(sym) >= 0):
        match = MANGLING_RE.search(i)
        if match is not None:
            sigString  =  match.group('signature')
            colonPos   =  sigString.rfind('::')
            pos        =  sigString.find(sym)
            symLen     =  pos + len(sym)
            if colonPos != -1 and colonPos != pos-2:   # Skips inexact substring matches like StartItem when search for Item
                continue
            if pos > 0 and sigString[symLen] == '(':   # Given a string 'funcName()...' - checks for '(' after name
                yield DecoratedSymbol(match.group('decorated'), sigString)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Utility for extracting mangled names'
                                                 ' from dumpbin\'s output for pragma mangle')
    parser.add_argument('file', help='D source file to scan')

    parser.add_argument('-o', '--out', type=str, help='optional output file name'
        ' (by default suffix `_gen` will be appended to source file name)')

    parser.add_argument('-e', '--exports', type=str, default='list_exports.txt',
                                           help='dumpbin exports file to read')
                                                         
    parser.add_argument('-p', '--prepend', type=str, help='file or custom text to prepend in generated file')
    args = parser.parse_args()
    prepend = args.prepend + os.linesep
    process_mangling(args.file, args.exports, args.out)