pypu

  1#!/usr/bin/env python3
  2"""
  3.. include:: ./example.svg
  4"""
  5
  6import argparse
  7import os
  8import sys
  9from argparse import _MutuallyExclusiveGroup
 10
 11import astroid
 12import requests
 13
 14from data.SourceFile import SourceFile
 15from data.UMLFile import UMLFile
 16from utils.ModuleParser import get_module_info
 17from utils.PlantEncoder import encode
 18
 19
 20def get_known_modules(directory: str, blacklist: list[str]) -> list[str]:
 21    """
 22    Gets a `list` of all modules in a directory and its subdirectories.
 23
 24    :param directory: The directory to process.
 25    :param blacklist: A `list` of forbidden substrings.
 26    :return: A `list` of known modules.
 27    """
 28    known_modules: list[str] = []
 29    for root, _, files in os.walk(directory):
 30        for file in files:
 31            if '.py' in file:
 32                file_path: str = os.path.join(root, file.split('.')[0])
 33                if not any(item in file_path for item in blacklist):
 34                    known_modules.append('.'.join([
 35                        x for x in file_path[len(directory):].split(os.sep) if x
 36                    ]))
 37    return known_modules
 38
 39
 40def get_source_files(
 41        directory: str,
 42        known_modules: list[str],
 43        blacklist: list[str]
 44) -> list[SourceFile]:
 45    """
 46    Gets a `list` of all source files in a directory and its subdirectories.
 47
 48    :param directory: The directory to process.
 49    :param known_modules: A `list` of known modules.
 50    :param blacklist: A `list` of forbidden substrings.
 51    :return: A `list` of source files.
 52    """
 53    source_files: list[SourceFile] = []
 54    for root, _, files in os.walk(directory):
 55        for file in files:
 56            if '.py' in file:
 57                file_path = os.path.join(root, file)
 58                if not any(item in file_path for item in blacklist):
 59                    try:
 60                        source_files.append(
 61                            get_module_info(directory, file_path, known_modules)
 62                        )
 63                    except astroid.AstroidBuildingException:
 64                        pass
 65    return source_files
 66
 67
 68def generate_uml(directory: str) -> str:
 69    """
 70    Generates a valid PlantUML string for all files in the given directory and its subdirectories.
 71
 72    :param directory: The directory to process.
 73    :return: A valid PlantUML string.
 74    """
 75    blacklist: list[str] = ['venv', '/.', '__']
 76    known_modules: list[str] = get_known_modules(directory, blacklist)
 77    source_files: list[SourceFile] = get_source_files(directory, known_modules, blacklist)
 78    return str(UMLFile(source_files))
 79
 80
 81if __name__ == '__main__':
 82    parser: argparse.ArgumentParser = argparse.ArgumentParser()
 83    parser.add_argument('-m', '--module', help='the module to analyze', required=True)
 84    group: _MutuallyExclusiveGroup = parser.add_mutually_exclusive_group()
 85    group.add_argument('-l', '--link', help='returns a link to the diagram in the specified format')
 86    group.add_argument('-f', '--format', help='returns the diagram in the specified format')
 87    args: argparse.Namespace = parser.parse_args()
 88
 89    uml: str = generate_uml(args.module)
 90    if args.link:
 91        print(encode(uml, args.link))
 92    elif args.format:
 93        encoded: str = encode(uml, args.format)
 94        result: requests.Response | None = None
 95        try:
 96            result = requests.get(encoded, timeout=3, allow_redirects=False)
 97        except requests.ReadTimeout:
 98            print(
 99                'The rendering server took too long to respond. Try visiting "' +
100                encoded +
101                '" in a browser.'
102            )
103        if result is not None and result.status_code == 200:
104            sys.stdout.buffer.write(result.content)
105    else:
106        print(uml)