From 723eeb73b53cacc1d55ff20f07e63ed6868701fb Mon Sep 17 00:00:00 2001 From: ed cuss Date: Wed, 24 Sep 2025 19:39:20 +0100 Subject: [PATCH 1/4] feat: save call tree as a json --- src/spaghettree/__main__.py | 35 +++++++++++++------------- src/spaghettree/adapters/io_wrapper.py | 16 ++++++++---- src/spaghettree/domain/entities.py | 10 +++++--- src/spaghettree/domain/processing.py | 5 +--- tests/test_main.py | 3 ++- 5 files changed, 39 insertions(+), 30 deletions(-) diff --git a/src/spaghettree/__main__.py b/src/spaghettree/__main__.py index 4733cab..715878c 100644 --- a/src/spaghettree/__main__.py +++ b/src/spaghettree/__main__.py @@ -1,11 +1,11 @@ import argparse -from pprint import pformat +import json +from pathlib import Path -from spaghettree import Ok, Result +from spaghettree import Result from spaghettree.adapters.io_wrapper import IOProtocol, IOWrapper from spaghettree.domain.adj_mat import AdjMat from spaghettree.domain.optimisation import ( - cyan, get_dwm, get_top_suggested_merges, yellow, @@ -39,27 +39,28 @@ def run_process( call_tree = entities_res.and_then(create_call_tree).unwrap() if optimise_src_code: - return optimise_entity_positions( - io=io, + res = optimise_entity_positions( entities=entities, location_map=location_map, call_tree=call_tree, src_root=src_root, new_root=new_root, + ).unwrap() + else: + adj_mat = AdjMat.from_call_tree_no_optimisation(call_tree).unwrap() + print( # noqa: T201 + yellow( + f"Current Directed Weighted Modularity (DWM): {get_dwm(adj_mat.mat, adj_mat.communities): .5f}" + ) ) + top_merges = get_top_suggested_merges(adj_mat).unwrap() - adj_mat = AdjMat.from_call_tree_no_optimisation(call_tree).unwrap() - print( # noqa: T201 - yellow( - f"Current Directed Weighted Modularity (DWM): {get_dwm(adj_mat.mat, adj_mat.communities): .5f}" - ) - ) - top_merges = get_top_suggested_merges(adj_mat).unwrap() + for merge in top_merges: + merge.display() - for merge in top_merges: - merge.display() + res = {Path("./call_tree.json").absolute(): json.dumps(call_tree, indent=4)} - return Ok(call_tree) + return io.write_files(res, ruff_root=new_root, format_code=optimise_src_code) if __name__ == "__main__": @@ -84,5 +85,5 @@ def run_process( args = parser.parse_args() res = main(args.src_root, new_root=args.new_root, optimise_src_code=args.optimise_src_code) - call_tree = res.unwrap() - print(f"\n{cyan(pformat(call_tree))}") # noqa: T201 + if not res.is_ok(): + raise res.error diff --git a/src/spaghettree/adapters/io_wrapper.py b/src/spaghettree/adapters/io_wrapper.py index 8ef849a..77bb3df 100644 --- a/src/spaghettree/adapters/io_wrapper.py +++ b/src/spaghettree/adapters/io_wrapper.py @@ -29,7 +29,9 @@ def read_files(self, root: str | Path) -> Result: ... @safe def write(self, modified_code: str, filepath: str, *, format_code: bool = True) -> None: ... - def write_files(self, src_code: dict[str, str], ruff_root: str | None = None) -> Result: ... + def write_files( + self, src_code: dict[str, str], ruff_root: str | None = None, *, format_code: bool = True + ) -> Result: ... @attrs.define @@ -70,11 +72,13 @@ def write(self, modified_code: str, filepath: str, *, format_code: bool = True) if format_code: self._run_ruff(filepath) - def write_files(self, src_code: dict[str, str], ruff_root: str | None = None) -> Result: + def write_files( + self, src_code: dict[str, str], ruff_root: str | None = None, *, format_code: bool = True + ) -> Result: results, fails = {}, {} for filepath, modified_code in src_code.items(): - if ruff_root is not None: + if not ruff_root or not format_code: # format all at the end instead res = self.write(modified_code, filepath, format_code=False) else: @@ -139,12 +143,14 @@ def read_files(self, root: str | Path) -> Result: def write(self, modified_code: str, filepath: str, *, format_code: bool = True) -> None: self.files[filepath] = format_code_str(modified_code) if format_code else modified_code - def write_files(self, src_code: dict[str, str], ruff_root: str | None = None) -> Result: + def write_files( + self, src_code: dict[str, str], ruff_root: str | None = None, *, format_code: bool = True + ) -> Result: results, fails = {}, {} for filepath, modified_code in src_code.items(): if ruff_root is not None: - res = self.write(modified_code, filepath) + res = self.write(modified_code, filepath, format_code=format_code) if res.is_ok(): results[filepath] = res.inner diff --git a/src/spaghettree/domain/entities.py b/src/spaghettree/domain/entities.py index 0847755..d09482b 100644 --- a/src/spaghettree/domain/entities.py +++ b/src/spaghettree/domain/entities.py @@ -136,9 +136,13 @@ def add_referenced_imports(self, imports: set[ImportCST]) -> Self: self.imports.update(imports) return self - for imp in imports: - if imp.as_name in self.referenced or imp.module in sys.stdlib_module_names: - self.imports.add(imp) + self.imports.update( + { + imp + for imp in imports + if imp.as_name in self.referenced or imp.module in sys.stdlib_module_names + } + ) return self diff --git a/src/spaghettree/domain/processing.py b/src/spaghettree/domain/processing.py index 323b559..bcd480f 100644 --- a/src/spaghettree/domain/processing.py +++ b/src/spaghettree/domain/processing.py @@ -4,7 +4,6 @@ from functools import partial from spaghettree import Result, safe -from spaghettree.adapters.io_wrapper import IOProtocol from spaghettree.domain.adj_mat import AdjMat from spaghettree.domain.entities import EntityCST, ImportCST, ImportType from spaghettree.domain.optimisation import ( @@ -19,8 +18,7 @@ from spaghettree.logger import logger -def optimise_entity_positions( # noqa: PLR0913 - io: IOProtocol, +def optimise_entity_positions( entities: dict[str, EntityCST], location_map: dict[str, EntityLocation], call_tree: dict[str, list[str]], @@ -44,7 +42,6 @@ def optimise_entity_positions( # noqa: PLR0913 ) .and_then(partial(create_new_filepaths, new_root=new_root or src_root)) .and_then(add_empty_inits_if_needed) - .and_then(partial(io.write_files, ruff_root=new_root or src_root)) ) diff --git a/tests/test_main.py b/tests/test_main.py index 27623cb..95b8096 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,3 +1,4 @@ +import json import os import shutil from pathlib import Path @@ -345,4 +346,4 @@ def test_run_process_return_call_tree(fixture_get_subset_files, expected_result) io = FakeIOWrapper(files) res = run_process(io, name, optimise_src_code=False) assert res.is_ok() - assert res.inner == expected_result + assert json.loads(io.files[Path("./call_tree.json").absolute()]) == expected_result From e1024bcac5ada2d1e43807fa939552f7d7a5d821 Mon Sep 17 00:00:00 2001 From: ed cuss Date: Wed, 24 Sep 2025 19:50:00 +0100 Subject: [PATCH 2/4] feat: expose args for call tree save path + docs --- README.md | 2 ++ src/spaghettree/__main__.py | 39 ++++++++++++++++++++++++++++++++----- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 3551775..48c5360 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,8 @@ Lastly it will print a representation of the call tree to the terminal to allow | ------------------ | --------------------- | -------- | ------- | ---------------------------------------------------- | | positional src_root | `str` | ✅ | | Path to the root of the repository to scan | | `--new-root` | `str` | ❌ | `''` | Optional new root path for output (default: empty, meaning same as src_root if optimisation is enabled). | +| `--call-tree-save-path` | `str` | ❌ | `'./call_tree.json'` | The location to save the generated call tree. Only used if `--optimise-src-code` isn't used. Defaults to `./call_tree.json`. | + | `--optimise-src-code` | Flag (no value) | ❌ | | Enable optimisation of the source code. | diff --git a/src/spaghettree/__main__.py b/src/spaghettree/__main__.py index 715878c..066801c 100644 --- a/src/spaghettree/__main__.py +++ b/src/spaghettree/__main__.py @@ -21,13 +21,30 @@ from spaghettree.logger import logger -def main(src_root: str, *, new_root: str = "", optimise_src_code: bool = False) -> Result: +def main( + src_root: str, + *, + new_root: str = "", + call_tree_save_path: str = "./call_tree.json", + optimise_src_code: bool = False, +) -> Result: io = IOWrapper() - return run_process(io, src_root, new_root=new_root, optimise_src_code=optimise_src_code) + return run_process( + io, + src_root, + new_root=new_root, + optimise_src_code=optimise_src_code, + call_tree_save_path=call_tree_save_path, + ) def run_process( - io: IOProtocol, src_root: str, *, new_root: str = "", optimise_src_code: bool = False + io: IOProtocol, + src_root: str, + *, + new_root: str = "", + call_tree_save_path: str = "./call_tree.json", + optimise_src_code: bool = False, ) -> Result: logger.info(f"*** RUNNING `spaghettree` {src_root = } {new_root = } ***") src_code = io.read_files(src_root).unwrap() @@ -58,7 +75,7 @@ def run_process( for merge in top_merges: merge.display() - res = {Path("./call_tree.json").absolute(): json.dumps(call_tree, indent=4)} + res = {Path(call_tree_save_path).absolute(): json.dumps(call_tree, indent=4)} return io.write_files(res, ruff_root=new_root, format_code=optimise_src_code) @@ -76,6 +93,13 @@ def run_process( default="", help="Optional new root path for output (default: empty, meaning same as src_root if optimisation is enabled).", ) + parser.add_argument( + "--call-tree-save-path", + dest="call_tree_save_path", + type=str, + default="./call_tree.json", + help="The location to save the generated call tree. Only used if `--optimise-src-code` isn't used. Defaults to `./call_tree.json`.", + ) parser.add_argument( "--optimise-src-code", dest="optimise_src_code", @@ -84,6 +108,11 @@ def run_process( ) args = parser.parse_args() - res = main(args.src_root, new_root=args.new_root, optimise_src_code=args.optimise_src_code) + res = main( + args.src_root, + new_root=args.new_root, + call_tree_save_path=args.call_tree_save_path, + optimise_src_code=args.optimise_src_code, + ) if not res.is_ok(): raise res.error From dc2e0d5439fef58ee090f4cf1af2f866e06bd146 Mon Sep 17 00:00:00 2001 From: ed cuss Date: Wed, 24 Sep 2025 20:01:26 +0100 Subject: [PATCH 3/4] fix: print where files are saved --- src/spaghettree/adapters/io_wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/spaghettree/adapters/io_wrapper.py b/src/spaghettree/adapters/io_wrapper.py index 77bb3df..6f45e18 100644 --- a/src/spaghettree/adapters/io_wrapper.py +++ b/src/spaghettree/adapters/io_wrapper.py @@ -12,6 +12,7 @@ from ruff.__main__ import find_ruff_bin from spaghettree import Err, Ok, Result, safe +from spaghettree.domain.optimisation import yellow from spaghettree.logger import logger @@ -86,6 +87,7 @@ def write_files( logger.debug(f"{filepath = } {res = }") if res.is_ok(): + print(yellow(f"File written to {filepath}")) # noqa: T201 results[filepath] = res.inner else: fails[filepath] = res From 6b6c428fd08b9029f4245891f6fa376be510d36a Mon Sep 17 00:00:00 2001 From: ed cuss Date: Wed, 24 Sep 2025 20:03:56 +0100 Subject: [PATCH 4/4] fix: print where files are saved --- src/spaghettree/__main__.py | 2 ++ src/spaghettree/adapters/io_wrapper.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/spaghettree/__main__.py b/src/spaghettree/__main__.py index 066801c..f1a0cea 100644 --- a/src/spaghettree/__main__.py +++ b/src/spaghettree/__main__.py @@ -64,6 +64,8 @@ def run_process( new_root=new_root, ).unwrap() else: + # remove any new_root so that it doesn't try to use ruff on the json + new_root = "" adj_mat = AdjMat.from_call_tree_no_optimisation(call_tree).unwrap() print( # noqa: T201 yellow( diff --git a/src/spaghettree/adapters/io_wrapper.py b/src/spaghettree/adapters/io_wrapper.py index 6f45e18..1204181 100644 --- a/src/spaghettree/adapters/io_wrapper.py +++ b/src/spaghettree/adapters/io_wrapper.py @@ -87,7 +87,7 @@ def write_files( logger.debug(f"{filepath = } {res = }") if res.is_ok(): - print(yellow(f"File written to {filepath}")) # noqa: T201 + print(yellow(f"File written to `{filepath}`")) # noqa: T201 results[filepath] = res.inner else: fails[filepath] = res