pretty-mod 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/CLAUDE.md +9 -6
  2. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/Cargo.lock +1 -1
  3. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/Cargo.toml +1 -1
  4. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/PKG-INFO +7 -5
  5. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/README.md +6 -4
  6. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/RELEASE_NOTES.md +29 -0
  7. pretty_mod-0.2.1/python/pretty_mod/__main__.py +6 -0
  8. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/python/pretty_mod/_pretty_mod.pyi +7 -2
  9. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/python/pretty_mod/cli.py +21 -2
  10. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/explorer.rs +2 -1
  11. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/lib.rs +28 -10
  12. pretty_mod-0.2.1/src/output_format.rs +121 -0
  13. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/signature.rs +26 -10
  14. pretty_mod-0.2.1/tests/test_json_output.py +103 -0
  15. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/.github/workflows/CI.yml +0 -0
  16. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/.github/workflows/tests.yml +0 -0
  17. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/.gitignore +0 -0
  18. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/.pre-commit-config.yaml +0 -0
  19. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/LICENSE +0 -0
  20. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/examples/0_hello.py +0 -0
  21. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/examples/1_tree.py +0 -0
  22. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/examples/2_sig.py +0 -0
  23. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/justfile +0 -0
  24. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/pyproject.toml +0 -0
  25. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/python/pretty_mod/__init__.py +0 -0
  26. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/python/pretty_mod/explorer.py +0 -0
  27. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/python/pretty_mod/py.typed +0 -0
  28. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/scripts/compare_local.py +0 -0
  29. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/scripts/compare_versions.py +0 -0
  30. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/scripts/perf_test.py +0 -0
  31. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/scripts/profile.py +0 -0
  32. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/config.rs +0 -0
  33. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/module_info.rs +0 -0
  34. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/package_downloader.rs +0 -0
  35. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/stdlib.rs +0 -0
  36. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/tree_formatter.rs +0 -0
  37. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/src/utils.rs +0 -0
  38. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/tests/__init__.py +0 -0
  39. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/tests/conftest.py +0 -0
  40. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/tests/test_cli.py +0 -0
  41. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/tests/test_double_colon.py +0 -0
  42. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/tests/test_explorer.py +0 -0
  43. {pretty_mod-0.2.0 → pretty_mod-0.2.1}/uv.lock +0 -0
@@ -2,30 +2,33 @@
2
2
 
3
3
  `pretty-mod` is a python package built on pyo3 to explore python packages for LLMs
4
4
 
5
- # Getting oriented
5
+ ## getting oriented
6
6
 
7
7
  - read @RELEASE_NOTES.md, @README.md, @pyproject.toml, and @justfile
8
8
 
9
- # run the tests
9
+ ## run the tests
10
10
 
11
11
  ```
12
12
  just test
13
13
  ```
14
14
 
15
- if for some reason you need to only build (just test does this automatically)
15
+ if you only need to build (`just test` runs `just build` automatically)
16
16
 
17
17
  ```
18
18
  just build
19
19
  ```
20
20
 
21
- # run the local python package
21
+ ## run the local python package
22
22
 
23
23
  ```
24
24
  uv run pretty-mod tree fastapi.routing
25
25
  ```
26
26
 
27
- # run the remote python package
27
+ ## run the remote python package
28
28
 
29
29
  ```
30
30
  uvx pretty-mod tree fastapi.routing
31
- ```
31
+ ```
32
+
33
+ # IMPORTANT
34
+ - avoid breaking changes to the public api defined by type stubs in @python/pretty_mod/_pretty_mod.pyi
@@ -1017,7 +1017,7 @@ dependencies = [
1017
1017
 
1018
1018
  [[package]]
1019
1019
  name = "pretty-mod"
1020
- version = "0.2.0"
1020
+ version = "0.2.1"
1021
1021
  dependencies = [
1022
1022
  "flate2",
1023
1023
  "pyo3",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "pretty-mod"
3
- version = "0.2.0"
3
+ version = "0.2.1"
4
4
  edition = "2021"
5
5
 
6
6
  [lib]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pretty-mod
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  License-File: LICENSE
5
5
  Summary: A python module tree explorer for LLMs (and humans)
6
6
  Author-email: zzstoatzz <thrast36@gmail.com>
@@ -12,11 +12,9 @@ Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
12
12
 
13
13
  a python module tree explorer for LLMs (and humans)
14
14
 
15
- > [!IMPORTANT]
16
- > for all versions `>=0.1.0`, wheels for different operating systems are built via `maturin` and published to pypi, install `<0.1.0` for a pure python version
17
-
18
15
  > [!NOTE]
19
- > Starting from v0.2.0, output includes colors by default. Use `PRETTY_MOD_NO_COLOR=1` to disable.
16
+ > - For all versions `>=0.1.0`, wheels for different operating systems are built via `maturin` and published to PyPI. Install `<0.1.0` for a pure Python version.
17
+ > - Starting from v0.2.0, output includes colors by default. Use `PRETTY_MOD_NO_COLOR=1` to disable.
20
18
 
21
19
  ```bash
22
20
  # Explore module structure
@@ -89,6 +87,10 @@ pretty-mod tree requests --depth 3
89
87
 
90
88
  # Display function signatures
91
89
  pretty-mod sig json:loads
90
+
91
+ # Get JSON output for programmatic use
92
+ pretty-mod tree json -o json | jq '.tree.submodules | keys'
93
+ pretty-mod sig json:dumps -o json | jq '.parameters'
92
94
  pretty-mod sig os.path:join
93
95
 
94
96
  # Explore packages even without having them installed
@@ -2,11 +2,9 @@
2
2
 
3
3
  a python module tree explorer for LLMs (and humans)
4
4
 
5
- > [!IMPORTANT]
6
- > for all versions `>=0.1.0`, wheels for different operating systems are built via `maturin` and published to pypi, install `<0.1.0` for a pure python version
7
-
8
5
  > [!NOTE]
9
- > Starting from v0.2.0, output includes colors by default. Use `PRETTY_MOD_NO_COLOR=1` to disable.
6
+ > - For all versions `>=0.1.0`, wheels for different operating systems are built via `maturin` and published to PyPI. Install `<0.1.0` for a pure Python version.
7
+ > - Starting from v0.2.0, output includes colors by default. Use `PRETTY_MOD_NO_COLOR=1` to disable.
10
8
 
11
9
  ```bash
12
10
  # Explore module structure
@@ -79,6 +77,10 @@ pretty-mod tree requests --depth 3
79
77
 
80
78
  # Display function signatures
81
79
  pretty-mod sig json:loads
80
+
81
+ # Get JSON output for programmatic use
82
+ pretty-mod tree json -o json | jq '.tree.submodules | keys'
83
+ pretty-mod sig json:dumps -o json | jq '.parameters'
82
84
  pretty-mod sig os.path:join
83
85
 
84
86
  # Explore packages even without having them installed
@@ -1,3 +1,28 @@
1
+ # Release Notes - v0.2.1
2
+
3
+ ## 📊 JSON Output Support & Better Type Annotation Handling
4
+
5
+ This release adds machine-readable JSON output and fixes a critical bug with complex type annotations.
6
+
7
+ ### ✨ New Features
8
+
9
+ - **📊 JSON Output Support**: Export tree and signature data as JSON for programmatic use
10
+ - `pretty-mod tree json -o json` - Get module structure as JSON
11
+ - `pretty-mod sig json:dumps -o json` - Get function signature as JSON
12
+ - Perfect for piping to `jq` or other JSON processors
13
+ - Follows the Kubernetes pattern of `-o <format>` for output selection
14
+ - Example: `pretty-mod tree json -o json | jq '.tree.submodules | keys'`
15
+
16
+ ### 🏗️ Technical Improvements
17
+
18
+ - **Visitor Pattern**: Implemented output formatters using the Visitor pattern for extensibility
19
+ - Clean separation between data structure and formatting
20
+ - Easy to add new output formats in the future
21
+ - Type-safe implementation using Rust traits
22
+
23
+
24
+ ---
25
+
1
26
  # Release Notes - v0.2.0
2
27
 
3
28
  ## 🎨 Customizable Display & Colors + Enhanced Signature Support
@@ -64,6 +89,10 @@ This release introduces customizable display characters, color output, full type
64
89
 
65
90
  ### 🐛 Bug Fixes
66
91
 
92
+ - **Complex type annotations**: Fixed parameter splitting for nested generics
93
+ - Previously: `Callable[[Any], str]` would split incorrectly on the comma
94
+ - Now: Properly handles all nested brackets and quotes in type annotations
95
+ - Affects all complex types like `Dict[str, List[int]]`, `Literal['a', 'b']`, etc.
67
96
  - **Stdlib module handling**: Built-in modules no longer trigger PyPI download attempts
68
97
  - **Signature discovery**: Improved recursive search for symbols exported in `__all__`
69
98
  - **Download messages**: Colored warning messages for better visibility
@@ -0,0 +1,6 @@
1
+ """Enable running pretty-mod as a module: python -m pretty_mod"""
2
+
3
+ from .cli import main
4
+
5
+ if __name__ == "__main__":
6
+ main()
@@ -14,7 +14,12 @@ class ModuleTreeExplorer:
14
14
  def get_tree_string(self) -> str: ...
15
15
 
16
16
  def display_tree(
17
- root_module_path: str, max_depth: int = 2, quiet: bool = False
17
+ root_module_path: str,
18
+ max_depth: int = 2,
19
+ quiet: bool = False,
20
+ format: str = "pretty",
18
21
  ) -> None: ...
19
- def display_signature(import_path: str, quiet: bool = False) -> str: ...
22
+ def display_signature(
23
+ import_path: str, quiet: bool = False, format: str = "pretty"
24
+ ) -> str: ...
20
25
  def import_object(import_path: str) -> Any: ...
@@ -26,6 +26,14 @@ def main():
26
26
  action="store_true",
27
27
  help="Suppress warnings and informational messages",
28
28
  )
29
+ tree_parser.add_argument(
30
+ "-o",
31
+ "--output",
32
+ type=str,
33
+ choices=["pretty", "json"],
34
+ default="pretty",
35
+ help="Output format (default: pretty)",
36
+ )
29
37
 
30
38
  sig_parser = subparsers.add_parser("sig", help="Display function signature")
31
39
  sig_parser.add_argument(
@@ -36,14 +44,25 @@ def main():
36
44
  action="store_true",
37
45
  help="Suppress download messages",
38
46
  )
47
+ sig_parser.add_argument(
48
+ "-o",
49
+ "--output",
50
+ type=str,
51
+ choices=["pretty", "json"],
52
+ default="pretty",
53
+ help="Output format (default: pretty)",
54
+ )
39
55
 
40
56
  args = parser.parse_args()
41
57
 
42
58
  try:
43
59
  if args.command == "tree":
44
- display_tree(args.module, args.depth, args.quiet)
60
+ # Call display_tree with format parameter
61
+ display_tree(args.module, args.depth, args.quiet, args.output)
45
62
  elif args.command == "sig":
46
- print(display_signature(args.import_path, args.quiet))
63
+ # Call display_signature with format parameter
64
+ result = display_signature(args.import_path, args.quiet, args.output)
65
+ print(result)
47
66
  else:
48
67
  parser.print_help()
49
68
  sys.exit(1)
@@ -1,4 +1,5 @@
1
1
  use crate::module_info::ModuleInfo;
2
+ use crate::tree_formatter::format_tree_display;
2
3
  use pyo3::prelude::*;
3
4
  use std::fs;
4
5
  use std::path::{Path, PathBuf};
@@ -116,7 +117,7 @@ impl ModuleTreeExplorer {
116
117
  };
117
118
 
118
119
  // Use the display_tree formatting logic, which expects the wrapped format
119
- crate::format_tree_display(py, &tree_obj, &self.root_module_path)
120
+ format_tree_display(py, &tree_obj, &self.root_module_path)
120
121
  }
121
122
  }
122
123
 
@@ -1,6 +1,7 @@
1
1
  mod config;
2
2
  mod explorer;
3
3
  mod module_info;
4
+ mod output_format;
4
5
  mod package_downloader;
5
6
  mod signature;
6
7
  mod stdlib;
@@ -8,15 +9,15 @@ mod tree_formatter;
8
9
  mod utils;
9
10
 
10
11
  use crate::explorer::ModuleTreeExplorer;
11
- use crate::signature::display_signature as display_signature_impl;
12
- use crate::tree_formatter::format_tree_display;
12
+ use crate::output_format::create_formatter;
13
13
  use crate::utils::{extract_base_package, try_download_and_import, import_object_impl};
14
14
  use pyo3::prelude::*;
15
15
 
16
16
  /// Display a module tree
17
17
  #[pyfunction]
18
- #[pyo3(signature = (root_module_path, max_depth = 2, quiet = false))]
19
- fn display_tree(py: Python, root_module_path: &str, max_depth: usize, quiet: bool) -> PyResult<()> {
18
+ #[pyo3(signature = (root_module_path, max_depth = 2, quiet = false, format = "pretty"))]
19
+ fn display_tree(py: Python, root_module_path: &str, max_depth: usize, quiet: bool, format: &str) -> PyResult<()> {
20
+ let formatter = create_formatter(format);
20
21
  // Check for invalid single colon (but allow double colon)
21
22
  if root_module_path.contains(':') && !root_module_path.contains("::") {
22
23
  return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
@@ -38,8 +39,8 @@ fn display_tree(py: Python, root_module_path: &str, max_depth: usize, quiet: boo
38
39
  let explorer = ModuleTreeExplorer::new(module_name.to_string(), max_depth);
39
40
  match explorer.explore(py) {
40
41
  Ok(tree) => {
41
- // Display tree using the wrapped format
42
- let tree_str = format_tree_display(py, &tree, module_name)?;
42
+ // Display tree using the formatter
43
+ let tree_str = formatter.format_tree(py, &tree, module_name)?;
43
44
  println!("{}", tree_str);
44
45
  Ok(())
45
46
  }
@@ -69,7 +70,7 @@ fn display_tree(py: Python, root_module_path: &str, max_depth: usize, quiet: boo
69
70
  let explorer = ModuleTreeExplorer::new(module_name.to_string(), max_depth);
70
71
  match explorer.explore(py) {
71
72
  Ok(tree) => {
72
- let tree_str = format_tree_display(py, &tree, module_name)?;
73
+ let tree_str = formatter.format_tree(py, &tree, module_name)?;
73
74
  println!("{}", tree_str);
74
75
  Ok(())
75
76
  }
@@ -109,9 +110,26 @@ fn display_tree(py: Python, root_module_path: &str, max_depth: usize, quiet: boo
109
110
 
110
111
  /// Display a function signature
111
112
  #[pyfunction]
112
- #[pyo3(signature = (import_path, quiet = false))]
113
- fn display_signature(py: Python, import_path: &str, quiet: bool) -> PyResult<String> {
114
- display_signature_impl(py, import_path, quiet)
113
+ #[pyo3(signature = (import_path, quiet = false, format = "pretty"))]
114
+ fn display_signature(py: Python, import_path: &str, quiet: bool, format: &str) -> PyResult<String> {
115
+ use crate::signature::try_ast_signature;
116
+ let formatter = create_formatter(format);
117
+
118
+ // First try to get signature from AST
119
+ if let Some(result) = try_ast_signature(py, import_path, quiet) {
120
+ if let Some(ref sig) = result.signature {
121
+ return Ok(formatter.format_signature(sig));
122
+ }
123
+ }
124
+
125
+ // If AST parsing didn't find it, return a simple message
126
+ let object_name = if import_path.contains(':') {
127
+ import_path.split(':').last().unwrap_or(import_path)
128
+ } else {
129
+ import_path.split('.').last().unwrap_or(import_path)
130
+ };
131
+
132
+ Ok(formatter.format_signature_not_available(object_name))
115
133
  }
116
134
 
117
135
  /// Import an object from a module path (public API, no auto-download)
@@ -0,0 +1,121 @@
1
+ use crate::module_info::FunctionSignature;
2
+ use pyo3::prelude::*;
3
+ use std::collections::HashMap;
4
+
5
+ /// Trait for different output format visitors
6
+ pub trait OutputFormatter {
7
+ /// Format a module tree
8
+ fn format_tree(&self, py: Python, tree: &PyObject, module_name: &str) -> PyResult<String>;
9
+
10
+ /// Format a function signature
11
+ fn format_signature(&self, signature: &FunctionSignature) -> String;
12
+
13
+ /// Format a signature not available message
14
+ fn format_signature_not_available(&self, object_name: &str) -> String;
15
+ }
16
+
17
+ /// Pretty print formatter (current default behavior)
18
+ pub struct PrettyPrintFormatter;
19
+
20
+ impl OutputFormatter for PrettyPrintFormatter {
21
+ fn format_tree(&self, py: Python, tree: &PyObject, module_name: &str) -> PyResult<String> {
22
+ // Use existing tree formatter
23
+ crate::tree_formatter::format_tree_display(py, tree, module_name)
24
+ }
25
+
26
+ fn format_signature(&self, signature: &FunctionSignature) -> String {
27
+ // Use existing signature formatter
28
+ crate::signature::format_signature_display(signature)
29
+ }
30
+
31
+ fn format_signature_not_available(&self, object_name: &str) -> String {
32
+ let config = crate::config::DisplayConfig::get();
33
+ format!(
34
+ "{} {} (signature not available)",
35
+ crate::config::colorize(
36
+ &config.signature_icon,
37
+ &config.color_scheme.signature_color,
38
+ config
39
+ ),
40
+ crate::config::colorize(object_name, &config.color_scheme.signature_color, config)
41
+ )
42
+ }
43
+ }
44
+
45
+ /// JSON formatter for machine-readable output
46
+ pub struct JsonFormatter;
47
+
48
+ impl OutputFormatter for JsonFormatter {
49
+ fn format_tree(&self, py: Python, tree: &PyObject, module_name: &str) -> PyResult<String> {
50
+ // Convert PyObject tree to a serializable structure
51
+ let mut result = HashMap::new();
52
+ result.insert(
53
+ "module".to_string(),
54
+ serde_json::Value::String(module_name.to_string()),
55
+ );
56
+
57
+ // Convert the tree structure to JSON
58
+ if let Ok(tree_value) = pyobject_to_json_value(py, tree) {
59
+ result.insert("tree".to_string(), tree_value);
60
+ }
61
+
62
+ serde_json::to_string_pretty(&result)
63
+ .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
64
+ }
65
+
66
+ fn format_signature(&self, signature: &FunctionSignature) -> String {
67
+ // Serialize signature to JSON
68
+ serde_json::to_string_pretty(signature).unwrap_or_else(|_| "{}".to_string())
69
+ }
70
+
71
+ fn format_signature_not_available(&self, object_name: &str) -> String {
72
+ let result = serde_json::json!({
73
+ "name": object_name,
74
+ "available": false,
75
+ "reason": "signature not available"
76
+ });
77
+ serde_json::to_string_pretty(&result).unwrap_or_else(|_| "{}".to_string())
78
+ }
79
+ }
80
+
81
+ /// Convert PyObject to serde_json::Value
82
+ fn pyobject_to_json_value(py: Python, obj: &PyObject) -> PyResult<serde_json::Value> {
83
+ // Try to extract as different Python types
84
+ if let Ok(dict) = obj.extract::<HashMap<String, PyObject>>(py) {
85
+ let mut map = serde_json::Map::new();
86
+ for (key, value) in dict {
87
+ if let Ok(json_value) = pyobject_to_json_value(py, &value) {
88
+ map.insert(key, json_value);
89
+ }
90
+ }
91
+ Ok(serde_json::Value::Object(map))
92
+ } else if let Ok(list) = obj.extract::<Vec<PyObject>>(py) {
93
+ let vec: Vec<serde_json::Value> = list
94
+ .iter()
95
+ .filter_map(|item| pyobject_to_json_value(py, item).ok())
96
+ .collect();
97
+ Ok(serde_json::Value::Array(vec))
98
+ } else if let Ok(s) = obj.extract::<String>(py) {
99
+ Ok(serde_json::Value::String(s))
100
+ } else if let Ok(b) = obj.extract::<bool>(py) {
101
+ Ok(serde_json::Value::Bool(b))
102
+ } else if let Ok(i) = obj.extract::<i64>(py) {
103
+ Ok(serde_json::Value::Number(serde_json::Number::from(i)))
104
+ } else if let Ok(f) = obj.extract::<f64>(py) {
105
+ if let Some(num) = serde_json::Number::from_f64(f) {
106
+ Ok(serde_json::Value::Number(num))
107
+ } else {
108
+ Ok(serde_json::Value::Null)
109
+ }
110
+ } else {
111
+ Ok(serde_json::Value::Null)
112
+ }
113
+ }
114
+
115
+ /// Factory function to create formatter based on format string
116
+ pub fn create_formatter(format: &str) -> Box<dyn OutputFormatter> {
117
+ match format.to_lowercase().as_str() {
118
+ "json" => Box::new(JsonFormatter),
119
+ _ => Box::new(PrettyPrintFormatter),
120
+ }
121
+ }
@@ -187,7 +187,7 @@ fn find_signature_recursive<'a>(
187
187
  }
188
188
 
189
189
  /// Format a signature for display
190
- fn format_signature_display(sig: &FunctionSignature) -> String {
190
+ pub fn format_signature_display(sig: &FunctionSignature) -> String {
191
191
  let config = DisplayConfig::get();
192
192
  let mut result = format!(
193
193
  "{} {}\n",
@@ -243,8 +243,15 @@ fn format_signature_display(sig: &FunctionSignature) -> String {
243
243
  result
244
244
  }
245
245
 
246
+ /// Result of signature discovery
247
+ pub struct SignatureResult {
248
+ pub signature: Option<FunctionSignature>,
249
+ #[allow(dead_code)]
250
+ pub formatted_output: String,
251
+ }
252
+
246
253
  /// Try to get signature from AST parsing
247
- fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<String> {
254
+ pub fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<SignatureResult> {
248
255
  // Parse the full specification first
249
256
  let (package_override, path_without_package, version) =
250
257
  crate::utils::parse_full_spec(import_path);
@@ -266,7 +273,7 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
266
273
  };
267
274
 
268
275
  // Helper function to try exploration and get signature
269
- let try_get_signature = |py: Python| -> Option<String> {
276
+ let try_get_signature = |py: Python| -> Option<FunctionSignature> {
270
277
  // For builtin modules (implemented in C), we can't extract signatures from filesystem
271
278
  if crate::stdlib::is_builtin_module(module_path) {
272
279
  return None;
@@ -276,7 +283,7 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
276
283
  let explorer = crate::explorer::ModuleTreeExplorer::new(module_path.to_string(), 2);
277
284
  if let Ok(module_info) = explorer.explore_module_pure_filesystem(py, module_path) {
278
285
  if let Some(sig) = module_info.signatures.get(object_name) {
279
- return Some(format_signature_display(sig));
286
+ return Some(sig.clone());
280
287
  }
281
288
 
282
289
  // Check if it's in __all__ and search recursively
@@ -284,7 +291,7 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
284
291
  if all_exports.contains(&object_name.to_string()) {
285
292
  // Use the recursive search function to find it anywhere in the tree
286
293
  if let Some(sig) = find_signature_recursive(&module_info, object_name) {
287
- return Some(format_signature_display(sig));
294
+ return Some(sig.clone());
288
295
  }
289
296
  }
290
297
  }
@@ -298,7 +305,7 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
298
305
  if let Ok(root_info) = explorer.explore_module_pure_filesystem(py, root_package) {
299
306
  // Search recursively for the object
300
307
  if let Some(sig) = find_signature_recursive(&root_info, object_name) {
301
- return Some(format_signature_display(sig));
308
+ return Some(sig.clone());
302
309
  }
303
310
  }
304
311
  }
@@ -308,7 +315,10 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
308
315
 
309
316
  // First try direct filesystem exploration
310
317
  if let Some(sig) = try_get_signature(py) {
311
- return Some(sig);
318
+ return Some(SignatureResult {
319
+ signature: Some(sig.clone()),
320
+ formatted_output: format_signature_display(&sig),
321
+ });
312
322
  }
313
323
 
314
324
  // Check if this is a stdlib module - if so, don't try to download
@@ -336,17 +346,23 @@ fn try_ast_signature(py: Python, import_path: &str, quiet: bool) -> Option<Strin
336
346
  download_result = try_get_signature(py);
337
347
  Ok(())
338
348
  }) {
339
- return download_result;
349
+ if let Some(sig) = download_result {
350
+ return Some(SignatureResult {
351
+ signature: Some(sig.clone()),
352
+ formatted_output: format_signature_display(&sig),
353
+ });
354
+ }
340
355
  }
341
356
 
342
357
  None
343
358
  }
344
359
 
345
360
  /// Display a function signature
361
+ #[allow(dead_code)]
346
362
  pub fn display_signature(py: Python, import_path: &str, quiet: bool) -> PyResult<String> {
347
363
  // First try to get signature from AST
348
- if let Some(ast_sig) = try_ast_signature(py, import_path, quiet) {
349
- return Ok(ast_sig);
364
+ if let Some(result) = try_ast_signature(py, import_path, quiet) {
365
+ return Ok(result.formatted_output);
350
366
  }
351
367
 
352
368
  // If AST parsing didn't find it, return a simple message
@@ -0,0 +1,103 @@
1
+ """Test JSON output format functionality."""
2
+
3
+ import json
4
+ import os
5
+ import subprocess
6
+ import sys
7
+
8
+
9
+ def test_tree_json_output():
10
+ """Test tree command with JSON output."""
11
+ result = subprocess.run(
12
+ [sys.executable, "-m", "pretty_mod", "tree", "json", "-o", "json"],
13
+ capture_output=True,
14
+ text=True,
15
+ )
16
+
17
+ assert result.returncode == 0
18
+
19
+ # Parse JSON output
20
+ data = json.loads(result.stdout)
21
+
22
+ # Check structure
23
+ assert "module" in data
24
+ assert data["module"] == "json"
25
+ assert "tree" in data
26
+ assert "api" in data["tree"]
27
+ assert "submodules" in data["tree"]
28
+
29
+ # Check some expected content
30
+ api = data["tree"]["api"]
31
+ assert "dump" in api["functions"]
32
+ assert "loads" in api["functions"]
33
+ assert "JSONEncoder" in api["all"]
34
+
35
+
36
+ def test_signature_json_output():
37
+ """Test signature command with JSON output."""
38
+ result = subprocess.run(
39
+ [sys.executable, "-m", "pretty_mod", "sig", "json:dumps", "-o", "json"],
40
+ capture_output=True,
41
+ text=True,
42
+ )
43
+
44
+ assert result.returncode == 0
45
+
46
+ # Parse JSON output
47
+ data = json.loads(result.stdout)
48
+
49
+ # Check structure
50
+ assert "name" in data
51
+ assert data["name"] == "dumps"
52
+ assert "parameters" in data
53
+ assert "obj" in data["parameters"]
54
+ assert "skipkeys=False" in data["parameters"]
55
+ assert "return_type" in data
56
+
57
+
58
+ def test_signature_not_available_json():
59
+ """Test signature not available in JSON format."""
60
+ result = subprocess.run(
61
+ [sys.executable, "-m", "pretty_mod", "sig", "sys:maxsize", "-o", "json"],
62
+ capture_output=True,
63
+ text=True,
64
+ )
65
+
66
+ assert result.returncode == 0
67
+
68
+ # Parse JSON output
69
+ data = json.loads(result.stdout)
70
+
71
+ # Check structure
72
+ assert "name" in data
73
+ assert data["name"] == "maxsize"
74
+ assert "available" in data
75
+ assert data["available"] is False
76
+ assert "reason" in data
77
+
78
+
79
+ def test_default_output_unchanged():
80
+ """Test that default output (without -o flag) remains unchanged."""
81
+ # Test tree
82
+ result_tree = subprocess.run(
83
+ [sys.executable, "-m", "pretty_mod", "tree", "json"],
84
+ capture_output=True,
85
+ text=True,
86
+ env={**os.environ, "PRETTY_MOD_NO_COLOR": "1"},
87
+ )
88
+
89
+ assert result_tree.returncode == 0
90
+ assert "📦 json" in result_tree.stdout
91
+ assert "├── ⚡ functions:" in result_tree.stdout
92
+
93
+ # Test signature
94
+ result_sig = subprocess.run(
95
+ [sys.executable, "-m", "pretty_mod", "sig", "json:dumps"],
96
+ capture_output=True,
97
+ text=True,
98
+ env={**os.environ, "PRETTY_MOD_NO_COLOR": "1"},
99
+ )
100
+
101
+ assert result_sig.returncode == 0
102
+ assert "📎 dumps" in result_sig.stdout
103
+ assert "├── Parameters:" in result_sig.stdout
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes