mytree/
lib.rs

1use chrono::{DateTime, Local};
2use clap::{arg, Parser};
3use colored::*;
4use regex::Regex;
5use serde::Serialize;
6use std::collections::HashSet;
7use std::error::Error;
8use std::fmt::Debug;
9use std::io;
10use std::path::{Path, PathBuf};
11use std::time::SystemTime;
12use std::{fmt, fs};
13
14#[derive(Parser, Debug)]
15#[command(
16    author,
17    version,
18    about = "Mytree is a terminal tool to visualize your project structure.",
19    long_about = "You can use mytree to create custom visualizations of your project structure.\
20     The features supported are:\
21     1. Filtering results by file extensions\
22     2. Filtering results by regex matching\
23     3. Filtering results to include hidden files\
24     4. Enable long format output with file size and timestamps\
25     5. Sort results alphabetically (default)\
26     6. Sort results by file size\
27     7. Sort results by last updated timestamp\
28     8. Write results to a file as JSON
29     "
30)]
31pub struct Args {
32    #[arg(default_value = ".", help = "Root directory to start traversal")]
33    pub path: PathBuf,
34
35    #[arg(
36        short = 's',
37        long = "sort",
38        help = "Supply the argument with 'fs' to sort by file size, 'ts' to sort by last updated timestamp, or nothing to sort alphabetically (default)"
39    )]
40    pub sort_by: Option<String>,
41
42    #[arg(
43        short = 'e',
44        long = "extension",
45        help = "Filter by file extensions (e.g. -e rs -e toml)"
46    )]
47    pub extension_filters: Option<Vec<String>>,
48
49    #[arg(
50        short = 'a',
51        long = "all",
52        default_value_t = false,
53        help = "Include hidden files and directories"
54    )]
55    pub show_hidden: bool,
56
57    #[arg(
58        short = 'r',
59        long = "regex",
60        help = "Filter entries by matching name with regex"
61    )]
62    pub regex: Option<String>,
63
64    #[arg(
65        short = 'l',
66        long = "long",
67        default_value_t = false,
68        help = "Enable long format output with file size and timestamps"
69    )]
70    pub long_format: bool,
71
72    #[arg(
73        short = 'j',
74        long = "json",
75        help = "Write directory tree in JSON format"
76    )]
77    pub write_json: Option<String>,
78}
79
80struct PrintOptions {
81    sort_by: SortBy,
82    extension_filters: Option<HashSet<String>>,
83    show_hidden: bool,
84    regex_filter: Option<Regex>,
85    long_format: bool,
86    write_json: Option<String>,
87}
88
89struct Stats {
90    dirs: usize,
91    files: usize,
92    size: u64,
93}
94
95struct EntryMeta {
96    name: String,
97    path: PathBuf,
98    size: u64,
99    mtime: SystemTime,
100    is_dir: bool,
101}
102
103#[derive(Debug, Clone)]
104enum SortBy {
105    Alphabetical,
106    FileSize,
107    LastUpdatedTimestamp,
108}
109
110#[derive(Debug)]
111pub struct ArgParseError {
112    pub details: ArgParseErrorType,
113}
114
115#[derive(Debug)]
116pub enum ArgParseErrorType {
117    SortFlag(String),
118    BadExtension(String),
119    BadRegex(String),
120}
121
122impl fmt::Display for ArgParseErrorType {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        match self {
125            ArgParseErrorType::SortFlag(flag) => write!(
126                f,
127                "invalid sort flag \"{flag}\" (expected \"fs\" or \"ts\")"
128            ),
129            ArgParseErrorType::BadExtension(ext) => write!(f, "invalid extension \"{ext}\""),
130            ArgParseErrorType::BadRegex(msg) => write!(f, "invalid regex -> {msg}"),
131        }
132    }
133}
134
135impl fmt::Display for ArgParseError {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        write!(f, "argument error -> {}", self.details)
138    }
139}
140
141impl Error for ArgParseError {}
142
143#[derive(Debug)]
144pub struct TreeParseError {
145    pub details: TreeParseType,
146}
147
148#[derive(Debug)]
149pub enum TreeParseType {
150    Io(String),
151    InvalidInput(String),
152}
153
154impl fmt::Display for TreeParseType {
155    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
156        match self {
157            TreeParseType::Io(msg) => write!(f, "IO error -> {msg}"),
158            TreeParseType::InvalidInput(msg) => write!(f, "{msg}"),
159        }
160    }
161}
162
163impl fmt::Display for TreeParseError {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        write!(f, "{}", self.details)
166    }
167}
168
169impl Error for TreeParseError {}
170
171impl From<io::Error> for TreeParseError {
172    fn from(e: io::Error) -> Self {
173        TreeParseError {
174            details: TreeParseType::Io(e.to_string()),
175        }
176    }
177}
178
179#[derive(Debug)]
180pub enum ParseError {
181    Args(ArgParseError),
182    Tree(TreeParseError),
183}
184
185impl fmt::Display for ParseError {
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        match self {
188            ParseError::Args(e) => Debug::fmt(&e, f),
189            ParseError::Tree(e) => Debug::fmt(&e, f),
190        }
191    }
192}
193
194impl Error for ParseError {
195    fn source(&self) -> Option<&(dyn Error + 'static)> {
196        match self {
197            ParseError::Args(e) => Some(e),
198            ParseError::Tree(e) => Some(e),
199        }
200    }
201}
202
203impl From<ArgParseError> for ParseError {
204    fn from(e: ArgParseError) -> Self {
205        Self::Args(e)
206    }
207}
208impl From<TreeParseError> for ParseError {
209    fn from(e: TreeParseError) -> Self {
210        Self::Tree(e)
211    }
212}
213
214impl From<ParseError> for io::Error {
215    fn from(e: ParseError) -> io::Error {
216        io::Error::other(e)
217    }
218}
219
220#[derive(Debug, Serialize)]
221struct TreeNode {
222    name: String,
223    path: PathBuf,
224    size: u64,
225    mtime: SystemTime,
226    is_dir: bool,
227    children: Option<Vec<TreeNode>>,
228}
229
230fn create_print_options_from_args(args: Args) -> Result<PrintOptions, ParseError> {
231    let sort_by = match args.sort_by.as_deref() {
232        Some("fs") => SortBy::FileSize,
233        Some("ts") => SortBy::LastUpdatedTimestamp,
234        Some(bad) => {
235            return Err(ParseError::Args(ArgParseError {
236                details: ArgParseErrorType::SortFlag(bad.into()),
237            }));
238        }
239        None => SortBy::Alphabetical,
240    };
241
242    let extension_filters = if let Some(list) = args.extension_filters {
243        let mut set = HashSet::with_capacity(list.len());
244        for raw in list {
245            let ext = raw.trim_start_matches('.');
246            if ext.is_empty() {
247                return Err(ParseError::Args(ArgParseError {
248                    details: ArgParseErrorType::BadExtension(raw),
249                }));
250            }
251            set.insert(ext.to_ascii_lowercase());
252        }
253        Some(set)
254    } else {
255        None
256    };
257
258    let regex_filter = if let Some(pattern) = args.regex {
259        match Regex::new(&pattern) {
260            Ok(re) => Some(re),
261            Err(e) => {
262                return Err(ParseError::Args(ArgParseError {
263                    details: ArgParseErrorType::BadRegex(format!(
264                        "invalid regex \"{pattern}\": {e}"
265                    )),
266                }));
267            }
268        }
269    } else {
270        None
271    };
272
273    Ok(PrintOptions {
274        sort_by,
275        extension_filters,
276        show_hidden: args.show_hidden,
277        regex_filter,
278        long_format: args.long_format,
279        write_json: args.write_json,
280    })
281}
282
283/*
284Return a vector of ordered row-level entries at a point in the directory
285*/
286fn create_ordered_row_level_entries(
287    path: &Path,
288    opts: &PrintOptions,
289) -> Result<Vec<EntryMeta>, ParseError> {
290    let iter = fs::read_dir(path).map_err(|e| {
291        ParseError::Tree(TreeParseError {
292            details: TreeParseType::Io(format!("error reading directory {}: {e}", path.display())),
293        })
294    })?;
295
296    let mut meta_entries = Vec::new(); // allocate lazily
297
298    for dir_entry in iter {
299        let entry = dir_entry.map_err(|e| {
300            ParseError::Tree(TreeParseError {
301                details: TreeParseType::Io(format!(
302                    "error reading an entry in {}: {e}",
303                    path.display()
304                )),
305            })
306        })?;
307
308        let file_type = entry.file_type().map_err(|e| {
309            ParseError::Tree(TreeParseError {
310                details: TreeParseType::InvalidInput(format!(
311                    "could not determine file type for {}: {e}",
312                    entry.path().display()
313                )),
314            })
315        })?;
316
317        let name = entry.file_name().to_string_lossy().to_string();
318        let ext = entry
319            .path()
320            .extension()
321            .and_then(|s| s.to_str())
322            .unwrap_or("")
323            .to_ascii_lowercase();
324
325        if !opts.show_hidden && name.starts_with('.') {
326            continue;
327        }
328        if opts
329            .extension_filters
330            .as_ref()
331            .is_some_and(|set| !set.contains(ext.as_str()))
332        {
333            continue;
334        }
335        if opts
336            .regex_filter
337            .as_ref()
338            .is_some_and(|re| !re.is_match(&name))
339        {
340            continue;
341        }
342
343        let md = entry.metadata().map_err(|e| {
344            ParseError::Tree(TreeParseError {
345                details: TreeParseType::Io(format!(
346                    "failed to read metadata for {}: {e}",
347                    entry.path().display()
348                )),
349            })
350        })?;
351
352        meta_entries.push(EntryMeta {
353            name,
354            path: entry.path(),
355            size: md.len(),
356            mtime: md.modified().unwrap_or(SystemTime::UNIX_EPOCH),
357            is_dir: file_type.is_dir(),
358        });
359    }
360
361    Ok(sort_meta_entries(meta_entries, &opts.sort_by))
362}
363
364fn sort_meta_entries(mut meta_entries: Vec<EntryMeta>, sort_criteria: &SortBy) -> Vec<EntryMeta> {
365    match sort_criteria {
366        SortBy::Alphabetical => {
367            meta_entries.sort_by(|a, b| a.name.to_lowercase().cmp(&b.name.to_lowercase()));
368        }
369        SortBy::FileSize => {
370            meta_entries.sort_by(|a, b| a.size.cmp(&b.size));
371        }
372        SortBy::LastUpdatedTimestamp => {
373            meta_entries.sort_by(|a, b| a.mtime.cmp(&b.mtime));
374        }
375    }
376    meta_entries
377}
378
379/*
380Return a vector of ordered row-level entries at a point in the directory
381*/
382fn build_directory_tree(root_path: &Path, opts: &PrintOptions) -> Result<TreeNode, ParseError> {
383    let md = fs::metadata(root_path).map_err(|e| {
384        ParseError::Tree(TreeParseError {
385            details: TreeParseType::Io(format!(
386                "failed to read metadata for {}: {e}",
387                root_path.display()
388            )),
389        })
390    })?;
391
392    let entries = create_ordered_row_level_entries(root_path, opts)?;
393    let mut kids = Vec::with_capacity(entries.len());
394    for entry in entries {
395        kids.push(build_tree_node_from_entry_meta(entry, opts)?);
396    }
397
398    Ok(TreeNode {
399        name: root_path
400            .file_name()
401            .map(|s| s.to_string_lossy().into_owned())
402            .unwrap_or_else(|| root_path.display().to_string()),
403        path: root_path.to_owned(),
404        size: md.len(),
405        mtime: md.modified().unwrap_or(SystemTime::UNIX_EPOCH),
406        is_dir: true,
407        children: Some(kids),
408    })
409}
410
411fn build_tree_node_from_entry_meta(
412    entry: EntryMeta,
413    opts: &PrintOptions,
414) -> Result<TreeNode, ParseError> {
415    let children = if entry.is_dir {
416        let subs = create_ordered_row_level_entries(&entry.path, opts)?;
417        let mut nodes = Vec::with_capacity(subs.len());
418        for sub in subs {
419            nodes.push(build_tree_node_from_entry_meta(sub, opts)?);
420        }
421        Some(nodes)
422    } else {
423        None
424    };
425
426    Ok(TreeNode {
427        name: entry.name,
428        path: entry.path,
429        size: entry.size,
430        mtime: entry.mtime,
431        is_dir: entry.is_dir,
432        children,
433    })
434}
435
436/*
437Print the directory tree to standard out or write to JSON
438*/
439fn print_tree(
440    node: &TreeNode,
441    connector: &str,
442    prefix_continuation: &str,
443    stats: &mut Stats,
444    opts: &PrintOptions,
445    write_fn: &mut dyn FnMut(&str),
446) {
447    let line = format_entry_line(&node.path, &node.name, opts.long_format);
448    write_fn(&format!("{prefix_continuation}{connector}{line}"));
449
450    if node.is_dir {
451        stats.dirs += 1;
452    } else {
453        stats.files += 1;
454        stats.size += node.size;
455    }
456
457    if let Some(children) = node.children.as_ref() {
458        let last_pos = children.len().saturating_sub(1);
459
460        for (idx, child) in children.iter().enumerate() {
461            let is_last = idx == last_pos;
462            let child_conn = if is_last { "└── " } else { "├── " };
463            let new_prefix = if is_last {
464                format!("{prefix_continuation}    ")
465            } else {
466                format!("{prefix_continuation}│   ")
467            };
468
469            print_tree(child, child_conn, &new_prefix, stats, opts, write_fn);
470        }
471    }
472}
473
474fn print_ascii_tree(root: &TreeNode, opts: &PrintOptions, root_path: &Path) {
475    let mut stats = Stats {
476        dirs: 0,
477        files: 0,
478        size: 0,
479    };
480
481    println!("{}", root_path.display());
482
483    let mut push_line = |line: &str| println!("{line}");
484
485    if let Some(children) = root.children.as_ref() {
486        let last = children.len().saturating_sub(1);
487        for (idx, child) in children.iter().enumerate() {
488            let is_last = idx == last;
489            let connector = if is_last { "└── " } else { "├── " };
490            let prefix = if is_last { "    " } else { "│   " };
491
492            print_tree(child, connector, prefix, &mut stats, opts, &mut push_line);
493        }
494    }
495
496    println!(
497        "\n{} directories, {} files, {} bytes total",
498        stats.dirs,
499        stats.files,
500        format_size(stats.size)
501    );
502}
503
504fn format_entry_line(path: &Path, name: &str, long_format: bool) -> String {
505    let is_hidden = name.starts_with('.') && name != "." && name != "..";
506    let styled_name = if path.is_dir() {
507        if is_hidden {
508            name.blue().bold().dimmed().underline()
509        } else {
510            name.blue().bold()
511        }
512    } else if is_hidden {
513        name.dimmed().underline()
514    } else {
515        match path
516            .extension()
517            .and_then(|e| e.to_str())
518            .map(|e| e.to_lowercase())
519        {
520            Some(ext) if ext == "rs" => name.red().bold(),
521            Some(ext) if ext == "py" => name.yellow().bold(),
522            Some(ext) if ["c", "cpp", "h", "hpp"].contains(&ext.as_str()) => name.cyan().bold(),
523            Some(ext) if ext == "cs" => name.magenta().bold(),
524            Some(ext) if ext == "ml" || ext == "mli" => name.bright_green().bold(),
525            Some(ext) if ext == "md" => name.white().italic(),
526            Some(ext) if ext == "txt" => name.dimmed(),
527            Some(ext) if ext == "json" => name.bright_yellow().bold(),
528            _ => name.normal(),
529        }
530    };
531
532    if long_format {
533        match fs::metadata(path) {
534            Ok(metadata) => {
535                let size = format_size(metadata.len());
536                let modified = metadata
537                    .modified()
538                    .ok()
539                    .map(format_time)
540                    .unwrap_or_else(|| "-".to_string());
541                let created = metadata
542                    .created()
543                    .ok()
544                    .map(format_time)
545                    .unwrap_or_else(|| "-".to_string());
546                format!(
547                    "{}\n      {:<10} {:<12} {:<10} {:<20} {:<10} {:<20}",
548                    styled_name, "Size:", size, "Modified:", modified, "Created:", created
549                )
550            }
551            Err(e) => format!("{styled_name} (Error reading metadata: {e})"),
552        }
553    } else {
554        styled_name.to_string()
555    }
556}
557
558fn format_size(bytes: u64) -> String {
559    const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
560    let mut size = bytes as f64;
561    let mut i = 0;
562    while size >= 1024.0 && i < UNITS.len() - 1 {
563        size /= 1024.0;
564        i += 1;
565    }
566    format!("{:.1} {:<2}", size, UNITS[i])
567}
568
569fn format_time(system_time: SystemTime) -> String {
570    let datetime: DateTime<Local> = system_time.into();
571    datetime.format("%Y-%m-%d %H:%M:%S").to_string()
572}
573
574fn write_tree_json<P>(nodes: &[TreeNode], dest: Option<P>) -> Result<(), ParseError>
575where
576    P: AsRef<Path>,
577{
578    let json_bytes = serde_json::to_vec_pretty(nodes).map_err(|e| {
579        ParseError::Tree(TreeParseError {
580            details: TreeParseType::InvalidInput(format!("serialising JSON: {e}")),
581        })
582    })?;
583
584    let path: PathBuf = dest
585        .map(|p| p.as_ref().to_path_buf())
586        .unwrap_or_else(|| PathBuf::from("file.json"));
587
588    if let Some(parent) = path.parent() {
589        fs::create_dir_all(parent).map_err(|e| {
590            ParseError::Tree(TreeParseError {
591                details: TreeParseType::Io(format!("creating {parent:?}: {e}")),
592            })
593        })?;
594    }
595
596    fs::write(&path, json_bytes).map_err(|e| {
597        ParseError::Tree(TreeParseError {
598            details: TreeParseType::Io(format!("writing {path:?}: {e}")),
599        })
600    })
601}
602
603fn emit_json(tree: &TreeNode, dest_raw: &str) -> Result<(), ParseError> {
604    let dest: Option<&Path> = if dest_raw.trim().is_empty() {
605        None
606    } else {
607        Some(Path::new(dest_raw))
608    };
609
610    write_tree_json(std::slice::from_ref(tree), dest)?;
611
612    println!(
613        "Wrote directory tree to {}",
614        dest.map(|p| p.display().to_string())
615            .unwrap_or_else(|| "file.json".into())
616    );
617
618    Ok(())
619}
620pub fn run(args: Args) -> io::Result<()> {
621    let path = &args.path.clone();
622    let opts = create_print_options_from_args(args)?;
623    let tree = build_directory_tree(path, &opts)?;
624
625    if let Some(ref raw_dest) = opts.write_json {
626        emit_json(&tree, raw_dest)?;
627        return Ok(());
628    }
629
630    print_ascii_tree(&tree, &opts, path);
631    Ok(())
632}
mytree/lib.rs

mytree/
lib.rs
