pyfm.py

import argparse
import sys
from datetime import datetime
from typing import List, Set, Tuple

HEADER = \
    "# THIS FILE IS GENERATED BY python-file-merger (https://github.com/weyh/python-file-merger).\n" + \
    "# DO NOT EDIT IT DIRECTLY!\n" + \
    f"# DATE: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"


class File:
    shebang: str
    file_path: str
    imports: List[str]
    prog_lines: List[str]

    @staticmethod
    def __parse_file(path: str) -> Tuple[str, List[str], List[str]]:
        shebang = ""
        imports: List[str] = []
        prog_lines: List[str] = []

        with open(path, "r") as f:
            for line in f:
                if line.startswith("#!"):
                    shebang = line
                elif line.startswith("import") or line.startswith("from"):
                    # move all imports to the top
                    imports.append(line.strip())
                else:
                    prog_lines.append(line)

        return (shebang, imports, prog_lines)

    def __init__(self, file_path: str):
        if not file_path.endswith(".py"):
            raise Exception("File must be a python file")

        self.file_path = file_path.replace("\\", "/")
        self.shebang, self.imports, self.prog_lines = File.__parse_file(file_path)

    def get_file_name(self) -> str:
        return self.file_path.split("/")[-1].replace(".py", "")


def create_output_file(file_path: str, files: List[File]):
    shebang = ""
    inline_import: List[str] = []
    imports: Set[str] = set()
    lines: List[str] = []

    # get shebang
    for file in files:
        if file.shebang != "":
            if shebang != "":
                raise Exception("Multiple shebangs found")
            else:
                shebang = file.shebang

    # get imports
    for file in files:
        for imp in file.imports:
            if any([file.get_file_name() in imp for file in files]):
                if imp.startswith("import"):
                    inline_import.append(file.get_file_name())
            else:
                imports.add(imp)

    # get program lines
    for file in files:
        lines.append(f"\n# {file.file_path}\n")

        for line in file.prog_lines:
            found = False
            for ii in inline_import:
                if f"{ii}." in line:
                    lines.append(line.replace(f"{ii}.", ""))
                    found = True

            if not found:
                lines.append(line)

    with open(file_path, "w") as f:
        if shebang != "":
            f.write(f"{shebang}\n")
        f.write(f"\n{HEADER}\n")
        f.writelines([f"{i}\n" for i in imports])
        f.writelines(lines)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("files", nargs="+", help="Files to merge")
    parser.add_argument("-o", "--output", help="Output file")

    return parser.parse_args()


def main() -> int:
    ret = 0
    pargs = parse_args()

    try:
        files = [File(path) for path in pargs.files]
        create_output_file(pargs.output, files)
    except Exception as e:
        print(e)
        ret = 1

    return ret


if __name__ == '__main__':
    sys.exit(main())