code-node-graph 0.1.0__py3-none-any.whl

code_node/cli.py

@@ -0,0 +1,316 @@
+"""
+code-node: run it in a folder of Python code and it opens a page showing how the code works from an entry point,
+as a graph of who calls whom and what each path reaches, and the steps of each call.
+
+    code-node                       # this folder: picks an entry point (asks when there are several)
+    code-node src/app/main.py       # start from a file: its main(), or pick one of its functions
+    code-node src/app/main.py:run   # ...or name the function (or Class.method) in it
+    code-node path/to/repo -e pkg.cli.main
+    code-node --list                # the entry points it found
+    code-node --all                 # every function in the graph, however deep
+    code-node --serve               # a live page on localhost that redraws whenever the code changes
+    code-node --editor cursor       # line links open in Cursor (vscode by default, none for no links)
+    code-node --plugins             # the platforms, libraries and deploy files it reads code with
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+import shutil
+import subprocess
+import sys
+import webbrowser
+from pathlib import Path
+
+from .deploys import deploys
+from .libraries import libraries
+from .pages import GRAPH_DEPTH, GRAPH_FUNCTIONS, PageOptions, make_mermaid, make_page
+from .platforms import platforms
+from .report.editor import EDITORS
+from .project import (
+    EntryPoint,
+    Project,
+    entries_in_file,
+    entry_in_files,
+    entry_points,
+    functions_matching,
+    project_root,
+    read_project,
+)
+
+_CACHE = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "code-node"
+
+
+def _parse_args(argv: list[str] | None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="code-node", description="See how Python code works: a graph of what calls what and what it reaches."
+    )
+    parser.add_argument(
+        "path",
+        nargs="?",
+        default=".",
+        help="the folder to read (default: this one), or the file to start from: main.py or main.py:function",
+    )
+    parser.add_argument(
+        "-e",
+        "--entry",
+        help="where to start: module.function, module.Class.method, "
+        "path/to/file.py[:function], or part of a function's name",
+    )
+    parser.add_argument("-l", "--list", action="store_true", help="list the entry points found, and stop")
+    parser.add_argument(
+        "--plugins",
+        action="store_true",
+        help="list the plugins it reads code with (platforms, libraries, deploy files), and stop",
+    )
+    parser.add_argument("-o", "--out", help="where to save the page (default: ~/.cache/code-node/)")
+    parser.add_argument("--no-open", action="store_true", help="save the page without opening a browser")
+    parser.add_argument("--all", action="store_true", help="draw every function, however deep (big graphs)")
+    parser.add_argument("--depth", type=int, default=GRAPH_DEPTH, help="calls deep the graph draws (default 5)")
+    parser.add_argument(
+        "--functions", type=int, default=GRAPH_FUNCTIONS, help="functions the graph draws at most (default 24)"
+    )
+    parser.add_argument("--mermaid", action="store_true", help="print the graph as Mermaid instead of a page")
+    parser.add_argument(
+        "--serve", action="store_true", help="serve the page on localhost and redraw it whenever a .py file changes"
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=None,
+        help="the --serve port (default: the first free one from 8765; 0 picks any free one)",
+    )
+    names = [p.name for p in platforms()]
+    parser.add_argument(
+        "--platform",
+        choices=names,
+        default="",
+        help="where the code runs, for that platform's checks (found entry points set it: a Lambda "
+        "handler is lambda; default plain Python)",
+    )
+    parser.add_argument(
+        "--lambda", dest="platform", action="store_const", const="lambda", help="the same as --platform lambda"
+    )
+    parser.add_argument(
+        "--editor",
+        choices=[*EDITORS, "none"],
+        default="vscode",
+        help="the editor the page's line links open (default vscode; none for no links)",
+    )
+    return parser.parse_args(argv)
+
+
+def _ask(candidates: list[EntryPoint]) -> EntryPoint:
+    """
+    The entry point the user picks from a numbered list (the first, when nobody's there to ask).
+    """
+    if len(candidates) == 1 or not sys.stdin.isatty():
+        return candidates[0]
+    _print_entries(candidates)
+    while True:
+        answer = input(f"Start from which one? [1-{len(candidates)}, Enter for 1] ").strip()
+        if not answer:
+            return candidates[0]
+        if answer.isdigit() and 1 <= int(answer) <= len(candidates):
+            return candidates[int(answer) - 1]
+
+
+def _print_entries(candidates: list[EntryPoint]) -> None:
+    width = max(len(c.name) for c in candidates)
+    for i, entry in enumerate(candidates, 1):
+        print(f"  {i:>2}. {entry.name:<{width}}  {entry.why}  ({entry.file}:{entry.line})")
+
+
+def _pick_entry(project: Project, given: str | None) -> EntryPoint | None:
+    """
+    The entry point to start from: the one given (exactly, or found by part of its name), else one found.
+    """
+    if given:
+        named = entry_in_files(project, given)
+        if named is not None:
+            return named
+        exact = [
+            c for c in entry_points(project) + functions_matching(project, given.rsplit(".", 1)[-1]) if c.name == given
+        ]
+        if exact:
+            return exact[0]
+        if "." in given and not functions_matching(project, given):
+            return EntryPoint(given, "given", "", 0)
+        matches = functions_matching(project, given)
+        if not matches:
+            print(f"No function in {project.root} matches {given!r}.", file=sys.stderr)
+            return None
+        return _ask(matches[:30])
+    found = entry_points(project)
+    if not found:
+        kinds = [
+            "console script",
+            "main()",
+            *(p.entry_kind.split(":")[0].split(" (")[0].removeprefix("a ") for p in platforms() if p.entry_kind),
+        ]
+        print(
+            f"Found no entry point in {project.root}: no {', '.join(kinds[:-1])}, or {kinds[-1]}. Name one with -e "
+            "module.function.",
+            file=sys.stderr,
+        )
+        return None
+    return _ask(found[:30])
+
+
+def _print_plugins() -> None:
+    """
+    Every plugin installed, by kind: built-in ones and those other packages add.
+    """
+    kinds = [
+        (
+            "Platforms (one per entry point: where it runs)",
+            [(p.name, p.title, p.entry_kind or "") for p in platforms()],
+        ),
+        (
+            "Libraries (as many as the code imports: what it reaches)",
+            [(lib.name, lib.title, ", ".join(lib.modules)) for lib in libraries()],
+        ),
+        (
+            "Deploy files (each found in the folder: how it's deployed)",
+            [(d.name, d.title, ", ".join(d.files)) for d in deploys()],
+        ),
+    ]
+    for title, rows in kinds:
+        print(title)
+        width = max((len(name) for name, _, _ in rows), default=0)
+        for name, label, what in rows:
+            print(f"  {name:<{width}}  {label}" + (f"  ({what})" if what else ""))
+
+
+def _is_wsl() -> bool:
+    version = Path("/proc/version")
+    return version.exists() and "microsoft" in version.read_text(errors="ignore").lower()
+
+
+def _open(target: str) -> None:
+    """
+    Opens a saved page (a path) or the server (a URL) in the browser, from WSL too (where Python's webbrowser often
+    can't reach Windows').
+    """
+    is_url = target.startswith("http")
+    if _is_wsl():
+        if shutil.which("wslview"):
+            subprocess.run(["wslview", target], check=False)
+            return
+        if shutil.which("explorer.exe"):
+            where = (
+                target
+                if is_url
+                else subprocess.run(
+                    ["wslpath", "-w", target], capture_output=True, text=True, check=False
+                ).stdout.strip()
+            )
+            subprocess.run(["explorer.exe", where], check=False)
+            return
+    webbrowser.open(target if is_url else Path(target).as_uri())
+
+
+def _out_path(given: str | None, project: Project, entry: str) -> Path:
+    if given:
+        return Path(given).expanduser().resolve()
+    slug = re.sub(r"[^\w.-]+", "-", f"{project.root.name}--{entry}")
+    return _CACHE / f"{slug}.html"
+
+
+def _split_file(text: str) -> tuple[Path, str] | None:
+    """
+    'src/app.py:run' -> (src/app.py, 'run'); 'src/app.py' -> (src/app.py, ''); None when it isn't a .py file.
+    """
+    file, _, function = text.partition(".py:") if ".py:" in text else (text, "", "")
+    path = Path(file + (".py" if function or text.endswith(".py:") else "")).expanduser()
+    return (path.resolve(), function) if path.suffix == ".py" and path.is_file() else None
+
+
+def _from_file(project: Project, file: Path, function: str) -> EntryPoint | None:
+    """
+    The entry point in a file: the function given, else the best of the file's (asking when there are several).
+    """
+    rel = file.relative_to(project.root).as_posix()
+    if rel not in project.files:
+        print(f"{rel} is skipped (a virtual environment, cache or build folder).", file=sys.stderr)
+        return None
+    found = entries_in_file(project, rel, function)
+    if not found:
+        print(
+            f"{rel} has no functions to start from: code-node follows a function (a script's top-level code "
+            "isn't one yet).",
+            file=sys.stderr,
+        )
+        return None
+    return _ask(found[:30])
+
+
+def _root_for(path: Path, file: Path | None) -> Path:
+    """
+    The folder to read: the one given, unless a file to start from was given that isn't in it (or as the path
+    itself), when it's that file's project.
+    """
+    if file is None:
+        return path
+    folder = path.resolve()
+    if path.is_dir() and folder in file.parents:
+        return folder if path != Path(".") else project_root(file)
+    return project_root(file)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """
+    The code-node command.
+    """
+    args = _parse_args(argv)
+    if args.plugins:
+        _print_plugins()
+        return 0
+    in_file = _split_file(args.path) or (_split_file(args.entry) if args.entry else None)
+    root = _root_for(Path(args.path).expanduser(), in_file[0] if in_file else None)
+    if not root.is_dir():
+        print(f"{root} isn't a folder or a .py file.", file=sys.stderr)
+        return 2
+    project = read_project(root)
+    if not project.files:
+        print(f"No Python files under {project.root}.", file=sys.stderr)
+        return 1
+    if args.list:
+        found = entry_points(project)
+        print(f"{len(found)} entry point{'s' if len(found) != 1 else ''} in {project.root}:")
+        _print_entries(found or [EntryPoint("(none: name one with -e module.function)", "", "-", 0)])
+        return 0
+    entry = _from_file(project, *in_file) if in_file else _pick_entry(project, args.entry)
+    if entry is None:
+        return 1
+    options = PageOptions(
+        draw_all=args.all, functions=args.functions, depth=args.depth, platform=args.platform, editor=args.editor
+    )
+    if args.mermaid:
+        print(make_mermaid(project.root, entry, options))
+        return 0
+    if args.serve:
+        from .serve import serve
+
+        return serve(project.root, entry, options, port=args.port, open_page=None if args.no_open else _open)
+    page, cmap = make_page(project.root, entry, options)
+    out = _out_path(args.out, project, entry.name)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(page, "utf-8")
+    services = list(dict.fromkeys(call.service for call in cmap.library_calls))
+    print(
+        f"{entry.name}: {len(cmap.functions)} functions, {len(cmap.aws)} AWS calls"
+        + (f", {len(cmap.library_calls)} through libraries ({', '.join(services)})" if services else "")
+        + (f" ({cmap.note})" if cmap.note else "")
+        + (f"\nDeployed as {cmap.function.name} in {cmap.function.source}" if cmap.function is not None else "")
+        + f"\nSaved {out}"
+    )
+    if not args.no_open:
+        _open(str(out))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code_node/deploys/__init__.py

@@ -0,0 +1,321 @@
+"""
+Deploy files: what a project's own deploy configuration says about where each entry point runs. A deploy plugin
+reads one kind of file (a SAM template, serverless.yml) and returns a Deployment for each function it declares: its
+handler, the Python it runs on, its timeout and environment variables, and what its role may call. Pure, like
+analysis/: the folder's reader hands it the files' text.
+
+Deploy files stack with the rest: a handler they name is an entry point even when it isn't called handler(), and
+its platform's checks then run against what the file says (an environment variable the code needs that isn't set,
+an AWS call its role doesn't allow, a module that Python version removed).
+
+Another package adds one by subclassing Deploy and naming it under the `code_node.deploys` entry point group:
+
+    [project.entry-points."code_node.deploys"]
+    cdk = "my_package.cdk:CdkApp"
+
+or, from a script, with register_deploy(CdkApp()).
+"""
+
+from __future__ import annotations
+
+import posixpath
+import re
+from fnmatch import fnmatchcase
+from typing import Any, Iterable
+
+from ..analysis.model import Deployment
+
+ENTRY_POINT_GROUP = "code_node.deploys"
+
+_BUILT_IN = ("code_node.deploys.sam:SamTemplate", "code_node.deploys.serverless:ServerlessFile")
+
+
+class Deploy:
+    """
+    A kind of deploy file. A plugin sets `name`, `title` and `files` (the file names it reads, wherever they are in
+    the folder), and overrides deployments(); grant() says how to allow what a function's role is missing.
+    """
+
+    name = ""  # 'sam'
+    title = ""  # 'AWS SAM'
+    files: tuple[str, ...] = ()  # 'template.yaml', 'template.yml'
+
+    def deployments(self, files: dict[str, str], code: set[str]) -> list[Deployment]:
+        """
+        The functions its files declare, each with the Python file and function it starts from. files is {path: text}
+        for the files named like its own; code is every Python file in the folder, to find handlers in.
+        """
+        return []
+
+    def grant(self, deployment: Deployment, actions: list[str]) -> str:
+        """
+        How to give a function's role IAM actions it's missing, as an instruction: 'add dynamodb:PutItem to its
+        role's policy'.
+        """
+        return f"allow {', '.join(actions)} in its role's policy in {deployment.source}"
+
+
+_deploys: dict[str, Deploy] = {}
+
+_loaded = False
+
+
+def register_deploy(deploy: Deploy) -> None:
+    """
+    Adds a deploy plugin (or replaces the one with its name): its files are read from then on.
+    """
+    if not deploy.name or not deploy.files:
+        raise ValueError("a deploy plugin needs a name and the files it reads")
+    deploys()
+    _deploys[deploy.name] = deploy
+
+
+def _add(found: Deploy | type[Deploy]) -> None:
+    deploy = found() if isinstance(found, type) else found
+    if deploy.name and deploy.files:
+        _deploys.setdefault(deploy.name, deploy)
+
+
+def deploys() -> list[Deploy]:
+    """
+    Every deploy plugin: the built-in ones, then those installed packages add, loaded the first time it's asked.
+    """
+    global _loaded
+    if not _loaded:
+        _loaded = True
+        from importlib import import_module
+        from importlib.metadata import entry_points
+
+        for target in _BUILT_IN:
+            module, _, attr = target.partition(":")
+            _add(getattr(import_module(module), attr))
+        for point in entry_points(group=ENTRY_POINT_GROUP):
+            try:
+                _add(point.load())
+            except Exception:  # noqa: BLE001  a broken plugin leaves the others working
+                continue
+    return list(_deploys.values())
+
+
+def get_deploy(name: str) -> Deploy | None:
+    return next((deploy for deploy in deploys() if deploy.name == name), None)
+
+
+def deploy_file_names() -> set[str]:
+    """
+    The file names any deploy plugin reads: what the folder's reader keeps besides .py files.
+    """
+    return {name for deploy in deploys() for name in deploy.files}
+
+
+def all_deployments(files: dict[str, str], code: Iterable[str]) -> list[Deployment]:
+    """
+    Every function the folder's deploy files declare, from every deploy plugin. A file a plugin can't read is
+    skipped: deploy files are a help, never a reason to fail.
+    """
+    code = set(code)
+    found: list[Deployment] = []
+    for deploy in deploys():
+        mine = {path: text for path, text in files.items() if posixpath.basename(path) in deploy.files}
+        if not mine:
+            continue
+        try:
+            found += deploy.deployments(mine, code)
+        except Exception:  # noqa: BLE001  a file it can't read
+            continue
+    return found
+
+
+def find_deployment(files: dict[str, str], code: Iterable[str], file: str, function: str) -> Deployment | None:
+    """
+    The deployment that starts at function() in file, if a deploy file declares one.
+    """
+    return next((d for d in all_deployments(files, code) if d.file == file and d.function == function), None)
+
+
+def allowed(action: str, permissions: Iterable[str]) -> bool:
+    """
+    Whether an IAM action is allowed by a set of actions, with their wildcards: 's3:GetObject' by 's3:Get*'.
+    """
+    action = action.lower()
+    return any(fnmatchcase(action, granted.lower()) for granted in permissions)
+
+
+def resolve_handler(base: str, code_dir: str, handler: str, code: set[str]) -> tuple[str, str] | None:
+    """
+    The Python file and function a handler setting names, from the deploy file's folder and the code's folder in
+    it: ('', 'src/', 'cart.app.on_checkout') -> ('src/cart/app.py', 'on_checkout'). None when no file matches.
+    """
+    module, _, function = str(handler).rpartition(".")
+    if not module or not function:
+        return None
+    stem = posixpath.normpath(posixpath.join(base, code_dir or ".", module.replace(".", "/")))
+    for candidate in (f"{stem}.py", f"{stem}/__init__.py"):
+        if candidate in code:
+            return candidate, function
+    return None
+
+
+def line_of(text: str, key: str, indent: int | None = None) -> int:
+    """
+    The line a mapping key is on in a YAML file ('CartFunction' -> 14), 0 when it can't be found.
+    """
+    spaces = r"[ \t]*" if indent is None else " " * indent
+    found = re.search(rf"^{spaces}['\"]?{re.escape(key)}['\"]?\s*:", text, re.M)
+    return text.count("\n", 0, found.start()) + 1 if found else 0
+
+
+def load_yaml(text: str) -> Any:
+    """
+    A YAML (or JSON) deploy file as plain data. CloudFormation's short tags (!Ref, !GetAtt, !Sub) are read as the
+    long forms ({'Ref': ...}, {'Fn::GetAtt': ...}); serverless.yml's ${...} variables stay as text. None when
+    PyYAML isn't installed or the file doesn't parse.
+    """
+    try:
+        import yaml
+    except ImportError:
+        return None
+
+    class Loader(yaml.SafeLoader):
+        pass
+
+    def tagged(loader: Any, suffix: str, node: Any) -> Any:
+        if isinstance(node, yaml.ScalarNode):
+            value: Any = loader.construct_scalar(node)
+        elif isinstance(node, yaml.SequenceNode):
+            value = loader.construct_sequence(node, deep=True)
+        else:
+            value = loader.construct_mapping(node, deep=True)
+        if suffix == "Ref":
+            return {"Ref": value}
+        if suffix == "GetAtt" and isinstance(value, str):
+            value = value.split(".", 1)
+        return {f"Fn::{suffix}": value}
+
+    Loader.add_multi_constructor("!", tagged)
+    try:
+        return yaml.load(text, Loader=Loader)  # noqa: S506  a SafeLoader, with CloudFormation's tags
+    except yaml.YAMLError:
+        return None
+
+
+def as_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    return list(value) if isinstance(value, list) else [value]
+
+
+def statement_actions(document: Any) -> set[str] | None:
+    """
+    The actions a policy document's statements allow: {'s3:GetObject', 'dynamodb:*'}. None when it can't be told:
+    a Deny or NotAction (what's left can't be listed), or an action that isn't written out.
+    """
+    if isinstance(document, dict) and "Statement" in document:
+        document = document["Statement"]
+    actions: set[str] = set()
+    for statement in as_list(document):
+        if not isinstance(statement, dict) or "NotAction" in statement or statement.get("Effect") != "Allow":
+            return None
+        for action in as_list(statement.get("Action")):
+            if not isinstance(action, str):
+                return None
+            actions.add(action)
+    return actions
+
+
+MANAGED_POLICIES: dict[str, set[str]] = {  # AWS managed policies -> the actions they allow, for the ones Lambda
+    # functions are given most (a policy not here leaves what a role allows unknown)
+    "AWSLambdaBasicExecutionRole": {"logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"},
+    "AWSLambdaExecute": {"logs:*", "s3:GetObject", "s3:PutObject"},
+    "AWSLambdaSQSQueueExecutionRole": {
+        "sqs:ReceiveMessage",
+        "sqs:DeleteMessage",
+        "sqs:GetQueueAttributes",
+        "logs:CreateLogGroup",
+        "logs:CreateLogStream",
+        "logs:PutLogEvents",
+    },
+    "AWSLambdaDynamoDBExecutionRole": {
+        "dynamodb:DescribeStream",
+        "dynamodb:GetRecords",
+        "dynamodb:GetShardIterator",
+        "dynamodb:ListStreams",
+        "logs:CreateLogGroup",
+        "logs:CreateLogStream",
+        "logs:PutLogEvents",
+    },
+    "AWSLambdaKinesisExecutionRole": {
+        "kinesis:DescribeStream",
+        "kinesis:DescribeStreamSummary",
+        "kinesis:GetRecords",
+        "kinesis:GetShardIterator",
+        "kinesis:ListShards",
+        "kinesis:ListStreams",
+        "kinesis:SubscribeToShard",
+        "logs:CreateLogGroup",
+        "logs:CreateLogStream",
+        "logs:PutLogEvents",
+    },
+    "AWSLambdaVPCAccessExecutionRole": {
+        "ec2:CreateNetworkInterface",
+        "ec2:DescribeNetworkInterfaces",
+        "ec2:DescribeSubnets",
+        "ec2:DeleteNetworkInterface",
+        "ec2:AssignPrivateIpAddresses",
+        "ec2:UnassignPrivateIpAddresses",
+        "logs:CreateLogGroup",
+        "logs:CreateLogStream",
+        "logs:PutLogEvents",
+    },
+    "AWSXRayDaemonWriteAccess": {
+        "xray:PutTraceSegments",
+        "xray:PutTelemetryRecords",
+        "xray:GetSamplingRules",
+        "xray:GetSamplingTargets",
+        "xray:GetSamplingStatisticSummaries",
+    },
+    "AmazonS3ReadOnlyAccess": {
+        "s3:Get*",
+        "s3:List*",
+        "s3:Describe*",
+        "s3-object-lambda:Get*",
+        "s3-object-lambda:List*",
+    },
+    "AmazonS3FullAccess": {"s3:*", "s3-object-lambda:*"},
+    "AmazonDynamoDBReadOnlyAccess": {
+        "dynamodb:BatchGetItem",
+        "dynamodb:ConditionCheckItem",
+        "dynamodb:Describe*",
+        "dynamodb:GetItem",
+        "dynamodb:List*",
+        "dynamodb:PartiQLSelect",
+        "dynamodb:Query",
+        "dynamodb:Scan",
+        "dynamodb:GetRecords",
+        "dynamodb:GetShardIterator",
+    },
+    "AmazonDynamoDBFullAccess": {"dynamodb:*"},
+    "AmazonSQSFullAccess": {"sqs:*"},
+    "AmazonSNSFullAccess": {"sns:*"},
+    "AmazonKinesisFullAccess": {"kinesis:*"},
+    "AmazonEventBridgeFullAccess": {"events:*"},
+    "AWSStepFunctionsFullAccess": {"states:*"},
+    "SecretsManagerReadWrite": {"secretsmanager:*"},
+    "AmazonSSMReadOnlyAccess": {"ssm:Describe*", "ssm:Get*", "ssm:List*"},
+    "AmazonBedrockFullAccess": {"bedrock:*"},
+    "AdministratorAccess": {"*"},
+}
+
+
+__all__ = [
+    "ENTRY_POINT_GROUP",
+    "MANAGED_POLICIES",
+    "Deploy",
+    "all_deployments",
+    "allowed",
+    "deploy_file_names",
+    "deploys",
+    "find_deployment",
+    "get_deploy",
+    "register_deploy",
+]