mod cli;
use clap::Parser as ClapParser;
use cli::{Cli, OutputFormat};
use color_eyre::Result;
use std::collections::HashMap;
use std::process;
use treemd::{Document, parser};
fn main() -> Result<()> {
color_eyre::install()?;
#[cfg(feature = "unstable-dynamic")]
clap_complete::CompleteEnv::with_factory(|| {
use clap::CommandFactory;
Cli::command()
})
.complete();
let args = Cli::parse();
#[cfg(feature = "unstable-dynamic")]
if args.setup_completions {
match cli::setup::setup_completions_interactive("treemd") {
Ok(_) => return Ok(()),
Err(e) => {
eprintln!("Error setting up completions: {}", e);
cli::setup::print_completion_instructions("treemd");
process::exit(1);
}
}
}
if args.query_help {
print_query_help();
return Ok(());
}
let (input_source, needs_file_picker, file_picker_dir) = match args.file.len() {
0 => {
use std::fs;
let cwd = std::env::current_dir().unwrap_or_default();
let md_files: Vec<_> = fs::read_dir(&cwd)
.ok()
.into_iter()
.flatten()
.filter_map(|entry| entry.ok())
.filter(|entry| {
let path = entry.path();
path.is_file()
&& path
.extension()
.and_then(|ext| ext.to_str())
.map(|ext| ext == "md" || ext == "markdown")
.unwrap_or(false)
})
.collect();
if md_files.is_empty() {
eprintln!("No markdown files found in current directory.");
eprintln!("\nUsage: treemd [OPTIONS] <FILE>");
eprintln!(" treemd [OPTIONS] -");
eprintln!(" treemd [OPTIONS] . # Open file picker");
eprintln!(" tree | treemd [OPTIONS]\n");
eprintln!("Tip: Navigate to a directory with .md files, or specify a file path.");
eprintln!("\nFor shell completion setup, use:");
eprintln!(" treemd --setup-completions");
std::process::exit(0);
}
(
treemd::input::InputSource::Stdin(
"# Select a file\n\nPress Enter to select a markdown file.".to_string(),
),
true,
None,
)
}
1 => {
let file_path = &args.file[0];
if file_path.is_dir() {
(
treemd::input::InputSource::Stdin(
"# Select a file\n\nPress Enter to select a markdown file.".to_string(),
),
true,
Some(file_path.clone()),
)
} else {
match treemd::input::determine_input_source(Some(file_path.as_path())) {
Ok(source) => (source, false, None),
Err(treemd::input::InputError::NoTty) => {
eprintln!("Error: markdown file argument is required");
eprintln!("\nUsage: treemd [OPTIONS] <FILE>");
eprintln!(" treemd [OPTIONS] -");
eprintln!(" treemd [OPTIONS] . # Open file picker");
eprintln!(" tree | treemd [OPTIONS]\n");
eprintln!(
"Use '-' to explicitly read from stdin, or pipe input with CLI flags."
);
eprintln!("\nFor shell completion setup, use:");
eprintln!(" treemd --setup-completions");
std::process::exit(1);
}
Err(e) => {
eprintln!("Error reading input: {}", e);
process::exit(1);
}
}
}
}
_ => {
let file_path = &args.file[0];
if file_path.is_dir() {
(
treemd::input::InputSource::Stdin(
"# Select a file\n\nPress Enter to select a markdown file.".to_string(),
),
true,
Some(file_path.clone()),
)
} else {
match treemd::input::determine_input_source(Some(file_path.as_path())) {
Ok(source) => (source, false, None),
Err(treemd::input::InputError::NoTty) => {
eprintln!("Error: markdown file argument is required");
std::process::exit(1);
}
Err(e) => {
eprintln!("Error reading input: {}", e);
process::exit(1);
}
}
}
}
};
let stdin_was_piped = matches!(input_source, treemd::input::InputSource::Stdin(_));
let markdown_content = match treemd::input::process_input(input_source) {
Ok(content) => content,
Err(e) => {
eprintln!("Error processing input: {}", e);
process::exit(1);
}
};
let doc = parser::parse_markdown(&markdown_content);
if let Some(ref query_str) = args.query {
return handle_query_mode(&doc, query_str, args.query_output.as_deref());
}
if !args.list
&& !args.tree
&& !args.count
&& args.section.is_none()
&& args.at_line.is_none()
&& !args.setup_completions
{
let mut config = treemd::Config::load();
if let Some(ref theme_name) = args.theme {
config.ui.theme = theme_name.clone();
}
let caps = treemd::tui::TerminalCapabilities::detect();
let color_mode = if let Some(ref mode_arg) = args.color_mode {
use cli::ColorModeArg;
use treemd::tui::ColorMode;
match mode_arg {
ColorModeArg::Auto => caps.recommended_color_mode,
ColorModeArg::Rgb => ColorMode::Rgb,
ColorModeArg::Color256 => ColorMode::Indexed256,
}
} else {
use treemd::tui::ColorMode;
match config.terminal.color_mode.as_str() {
"rgb" => ColorMode::Rgb,
"256" => ColorMode::Indexed256,
_ => caps.recommended_color_mode,
}
};
if caps.should_warn && !config.terminal.warned_terminal_app {
if let Some(warning) = caps.warning_message() {
eprintln!("\n{}\n", warning);
if !stdin_was_piped {
use std::io::{Read, stdin};
let _ = stdin().read(&mut [0u8]);
} else {
eprintln!("Press any key in the TUI to continue...");
}
}
let _ = config.set_warned_terminal_app();
}
use crossterm::ExecutableCommand;
use crossterm::event::EnableMouseCapture;
use crossterm::terminal::EnterAlternateScreen;
use std::io::stdout;
treemd::tui::tty::install_panic_hook();
treemd::tui::tty::enable_raw_mode().inspect_err(|e| {
eprintln!("Failed to enable raw mode: {}", e);
eprintln!("Note: When piping input, ensure you have a controlling terminal.");
})?;
stdout().execute(EnterAlternateScreen).inspect_err(|_| {
treemd::tui::tty::disable_raw_mode().ok();
})?;
let _ = stdout().execute(EnableMouseCapture);
let backend = ratatui::backend::CrosstermBackend::new(stdout());
let mut terminal = ratatui::Terminal::new(backend).inspect_err(|_| {
treemd::tui::tty::disable_raw_mode().ok();
})?;
let (filename, file_path) = if !args.file.is_empty() && !args.file[0].is_dir() {
let file = &args.file[0];
let name = file
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("stdin")
.to_string();
let path = file.canonicalize().unwrap_or_else(|_| file.clone());
(name, path)
} else {
("stdin".to_string(), std::path::PathBuf::from("<stdin>"))
};
let images_enabled = if args.no_images {
false
} else if args.images {
true
} else {
config.images.enabled
};
let mut app =
treemd::App::new(doc, filename, file_path, config, color_mode, images_enabled);
if needs_file_picker {
app.startup_needs_file_picker = true;
}
if let Some(dir) = file_picker_dir {
app.file_picker_dir = Some(dir.canonicalize().unwrap_or(dir));
}
let result = treemd::tui::run(&mut terminal, app);
use crossterm::event::DisableMouseCapture;
use crossterm::terminal::LeaveAlternateScreen;
stdout().execute(DisableMouseCapture).ok();
stdout().execute(LeaveAlternateScreen).ok();
treemd::tui::tty::disable_raw_mode().ok();
return result;
}
handle_cli_mode(&args, &doc);
Ok(())
}
fn handle_cli_mode(args: &Cli, doc: &Document) {
let headings: Vec<_> = if let Some(level) = args.level {
doc.headings_at_level(level)
} else if let Some(ref filter) = args.filter {
doc.filter_headings(filter)
} else {
doc.headings.iter().collect()
};
if let Some(line) = args.at_line {
print_heading_at_line(doc, line);
return;
}
if args.count {
print_heading_counts(doc);
} else if args.tree {
print_tree(doc, &args.output, &headings);
} else if let Some(ref section_name) = args.section {
extract_section(doc, section_name);
} else if args.list {
print_headings(&headings, &args.output, doc);
}
}
fn print_heading_at_line(doc: &Document, target_line: usize) {
if target_line == 0 {
eprintln!("Error: line number must be >= 1");
process::exit(1);
}
let mut best: Option<&parser::Heading> = None;
let mut prev_line = 1usize;
let mut prev_offset = 0usize;
for h in &doc.headings {
let h_line = prev_line + doc.content[prev_offset..h.offset].matches('\n').count();
if h_line <= target_line {
best = Some(h);
} else {
break;
}
prev_line = h_line;
prev_offset = h.offset;
}
match best {
Some(h) => {
let prefix = "#".repeat(h.level);
println!("{} {}", prefix, h.text);
}
None => {
eprintln!("No heading at or before line {}", target_line);
process::exit(1);
}
}
}
fn print_headings(headings: &[&parser::Heading], format: &OutputFormat, doc: &Document) {
match format {
OutputFormat::Plain => {
for heading in headings {
let prefix = "#".repeat(heading.level);
println!("{} {}", prefix, heading.text);
}
}
OutputFormat::Json => {
let json_output = parser::build_json_output(doc, None);
let json = serde_json::to_string_pretty(&json_output)
.expect("JSON serialization of document output should not fail");
println!("{}", json);
}
OutputFormat::Tree => {
eprintln!("Use --tree for tree output");
process::exit(1);
}
}
}
fn print_tree(doc: &Document, format: &OutputFormat, headings: &[&parser::Heading]) {
let tree = if headings.len() == doc.headings.len() {
doc.build_tree()
} else {
let owned: Vec<parser::Heading> = headings.iter().map(|&h| h.clone()).collect();
Document::new(String::new(), owned).build_tree()
};
let config = treemd::Config::load();
let compact = config.is_compact_tree();
match format {
OutputFormat::Tree | OutputFormat::Plain => {
for (i, node) in tree.iter().enumerate() {
let is_last = i == tree.len() - 1;
print!("{}", node.render_box_tree_styled("", is_last, compact));
}
}
OutputFormat::Json => {
let owned: Vec<parser::Heading> = headings.iter().map(|&h| h.clone()).collect();
let json = serde_json::to_string_pretty(&owned)
.expect("JSON serialization of headings should not fail");
println!("{}", json);
}
}
}
fn print_heading_counts(doc: &Document) {
let mut counts: HashMap<usize, usize> = HashMap::new();
for heading in &doc.headings {
*counts.entry(heading.level).or_insert(0) += 1;
}
println!("Heading counts:");
for level in 1..=6 {
if let Some(count) = counts.get(&level) {
let prefix = "#".repeat(level);
println!(" {}: {}", prefix, count);
}
}
println!("\nTotal: {}", doc.headings.len());
}
fn extract_section(doc: &Document, section_name: &str) {
let heading = match doc.find_heading(section_name) {
Some(h) => h,
None => {
eprintln!("Section '{}' not found", section_name);
process::exit(1);
}
};
let start = heading.offset;
let level = heading.level;
let end = doc
.headings
.iter()
.find(|h| h.offset > start && h.level <= level)
.map(|h| h.offset)
.unwrap_or(doc.content.len());
println!("{}", doc.content[start..end].trim());
}
fn handle_query_mode(doc: &Document, query_str: &str, output_format: Option<&str>) -> Result<()> {
use treemd::query::{self, OutputFormat};
let format = output_format
.map(|s| s.parse::<OutputFormat>())
.transpose()
.map_err(|e| {
eprintln!("Error: {}", e);
process::exit(1);
})?
.unwrap_or(OutputFormat::Plain);
match query::execute(doc, query_str) {
Ok(results) => {
if results.is_empty() {
return Ok(());
}
let output = query::format_output(&results, format);
println!("{}", output);
Ok(())
}
Err(e) => {
eprintln!("{}", e);
process::exit(1);
}
}
}
fn print_query_help() {
let help = r#"
treemd Query Language (tql)
A jq-like query language for navigating and extracting markdown structure.
ELEMENT SELECTORS
.h, .heading All headings (any level)
.h1 - .h6 Headings by level
.code All code blocks
.code[rust] Code blocks by language
.link, .a All links
.link[external] External links only
.img All images
.table All tables
.list All lists
.blockquote All blockquotes
FILTERS & INDEXING
.h2[Features] Heading containing "Features" (fuzzy)
.h2["Installation"] Heading with exact text
.h2[0] First h2
.h2[-1] Last h2
.h2[1:3] h2s at index 1 and 2
.h2[:3] First 3 h2s
HIERARCHY
.h1 > .h2 Direct child h2s under h1s
.h1 >> .code Code blocks anywhere under h1s
PIPES
.h2 | text Get heading text (strips ##)
[.h2] | count Count all h2s
.code | lang Get code block languages
.link | url Get link URLs
COLLECTION FUNCTIONS
count, length Count elements (alias: len, size)
first, last First/last element (alias: head)
limit(n), take(n) First n elements
skip(n), drop(n) Skip first n elements
nth(n) Get element at index
reverse Reverse order
sort Sort alphabetically
sort_by(key) Sort by property
unique Remove duplicates
flatten Flatten nested arrays
group_by(key) Group elements by key
min, max Min/max numeric value
add Sum numbers or concat strings
STRING FUNCTIONS
text Get text representation
upper, lower Case conversion
trim Strip whitespace
split(sep) Split by separator
join(sep) Join with separator
replace(a, b) Replace substring
slugify URL-friendly slug
lines, words, chars Count lines/words/chars
FILTER FUNCTIONS
select(cond) Keep if condition true (alias: where, filter)
contains(s) Contains substring (alias: includes)
startswith(s) Starts with prefix
endswith(s) Ends with suffix
matches(regex) Matches regex pattern
any, all Check if any/all truthy
not Negate boolean
CONTENT FUNCTIONS
content Section content (for headings)
md Raw markdown
url, href, src Get URL/link/image source
lang Code block language
AGGREGATION FUNCTIONS
stats Document statistics
levels Heading count by level
langs Code block count by language
types Link types count
EXAMPLES
# List all h2 headings
treemd -q '.h2' doc.md
# Get heading text only
treemd -q '.h2 | text' doc.md
# Count headings
treemd -q '[.h2] | count' doc.md
# First 5 headings
treemd -q '[.h] | limit(5)' doc.md
# Filter headings (three equivalent ways)
treemd -q '.h | select(contains("API"))' doc.md
treemd -q '.h | where(contains("API"))' doc.md
treemd -q '.h[API]' doc.md
# All Rust code blocks
treemd -q '.code[rust]' doc.md
# External link URLs
treemd -q '.link[external] | url' doc.md
# h2s under "Features" section
treemd -q '.h1[Features] > .h2' doc.md
# Group headings by level
treemd -q '[.h] | group_by("level")' doc.md
# Document statistics
treemd -q '. | stats' doc.md
# JSON output
treemd -q '.h2' --query-output json doc.md
OUTPUT FORMATS (--query-output)
plain Human-readable text (default)
json Compact JSON
json-pretty Pretty-printed JSON (alias: jsonp)
jsonl Line-delimited JSON (one per line)
md Raw markdown
tree Tree structure
For more details, see: https://github.com/epistates/treemd
"#;
println!("{}", help.trim());
}