Add firstLines handlebars helper

Cargo.toml

@@ -14,6 +14,7 @@ colored = "2.1.0"
 indicatif = "0.17.8"
 log = "0.4"
 num-format = { version = "0.4.4", features = ["with-system-locale"] }
+serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0.114"
 handlebars = "4.3"
 jwalk = "0.8"
@@ -35,3 +36,5 @@ env_logger = { version = "0.11.3" }
 arboard = { version = "3.4.0" }
 derive_builder = { version = "0.20.2" }
 winapi = { version = "0.3.9", features = ["errhandlingapi"] }
+unicode-width = "0.1"
+terminal_size = "0.3"

crates/code2prompt-core/Cargo.toml

@@ -27,6 +27,7 @@ colored = { workspace = true }
 indicatif = { workspace = true }
 log = { workspace = true }
 num-format = { workspace = true }
+serde = { workspace = true }
 serde_json = { workspace = true }
 handlebars = { workspace = true }
 jwalk = { workspace = true }

crates/code2prompt-core/build.rs

@@ -5,4 +5,4 @@ fn main() {
         println!("cargo:rustc-link-arg=-undefined");
         println!("cargo:rustc-link-arg=dynamic_lookup");
     }
-}
+}

crates/code2prompt-core/src/configuration.rs

@@ -101,6 +101,10 @@ pub struct Code2PromptConfig {
     /// Extra template data
     #[builder(default)]
     pub user_variables: HashMap<String, String>,
+
+    /// If true, token counting will be performed for each file (for token map display)
+    #[builder(default)]
+    pub token_map_enabled: bool,
 }
 
 impl Code2PromptConfig {

crates/code2prompt-core/src/filter.rs

@@ -26,7 +26,7 @@ pub fn build_globset(patterns: &[String]) -> GlobSet {
     } else {
         patterns.to_vec()
     };
-    
+
     for pattern in expanded_patterns {
         // If the pattern does not contain a '/' or the platform’s separator, prepend "**/"
         let normalized_pattern =
@@ -118,4 +118,4 @@ fn expand_brace_patterns(patterns: &[String]) -> Vec<String> {
         .split(',')
         .map(|expanded_pattern| format!("{}/{}", common_prefix, expanded_pattern))
         .collect::<Vec<String>>();
-}
+}

crates/code2prompt-core/src/path.rs

@@ -2,15 +2,32 @@
 use crate::configuration::Code2PromptConfig;
 use crate::filter::{build_globset, should_include_file};
 use crate::sort::{sort_files, sort_tree, FileSortMethod};
+use crate::tokenizer::count_tokens;
 use crate::util::strip_utf8_bom;
 use anyhow::Result;
 use ignore::WalkBuilder;
 use log::debug;
+use serde::{Deserialize, Serialize};
 use serde_json::json;
 use std::fs;
 use std::path::Path;
 use termtree::Tree;
 
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub struct EntryMetadata {
+    pub is_dir: bool,
+    pub is_symlink: bool,
+}
+
+impl From<&std::fs::Metadata> for EntryMetadata {
+    fn from(meta: &std::fs::Metadata) -> Self {
+        Self {
+            is_dir: meta.is_dir(),
+            is_symlink: meta.is_symlink(),
+        }
+    }
+}
+
 /// Traverses the directory and returns the string representation of the tree and the vector of JSON file representations.
 ///
 /// This function uses the provided configuration to determine which files to include, how to format them,
@@ -77,56 +94,72 @@ pub fn traverse_directory(config: &Code2PromptConfig) -> Result<(String, Vec<ser
 
             // ~~~ Processing File ~~~
             if path.is_file() && file_match {
-                if let Ok(code_bytes) = fs::read(path) {
-                    let clean_bytes = strip_utf8_bom(&code_bytes);
-                    let code = String::from_utf8_lossy(&clean_bytes);
-
-                    let code_block = wrap_code_block(
-                        &code,
-                        path.extension().and_then(|ext| ext.to_str()).unwrap_or(""),
-                        config.line_numbers,
-                        config.no_codeblock,
-                    );
-
-                    if !code.trim().is_empty() && !code.contains(char::REPLACEMENT_CHARACTER) {
-                        // ~~~ Filepath ~~~
-                        let file_path = if config.absolute_path {
-                            path.display().to_string()
-                        } else {
-                            format!("{}/{}", parent_directory, relative_path.display())
-                        };
-
-                        // ~~~ File JSON Representation ~~~
-                        let mut file_entry = serde_json::Map::new();
-                        file_entry.insert("path".to_string(), json!(file_path));
-                        file_entry.insert(
-                            "extension".to_string(),
-                            json!(path.extension().and_then(|ext| ext.to_str()).unwrap_or("")),
+                if let Ok(metadata) = entry.metadata() {
+                    if let Ok(code_bytes) = fs::read(path) {
+                        let clean_bytes = strip_utf8_bom(&code_bytes);
+                        let code = String::from_utf8_lossy(&clean_bytes);
+
+                        let code_block = wrap_code_block(
+                            &code,
+                            path.extension().and_then(|ext| ext.to_str()).unwrap_or(""),
+                            config.line_numbers,
+                            config.no_codeblock,
                         );
-                        file_entry.insert("code".to_string(), json!(code_block));
-
-                        // If date sorting is requested, record the file modification time.
-                        if let Some(method) = config.sort_method {
-                            if method == FileSortMethod::DateAsc
-                                || method == FileSortMethod::DateDesc
-                            {
-                                let mod_time = fs::metadata(path)
-                                    .and_then(|m| m.modified())
-                                    .and_then(|mtime| {
-                                        Ok(mtime.duration_since(std::time::SystemTime::UNIX_EPOCH))
-                                    })
-                                    .map(|d| d.unwrap().as_secs())
-                                    .unwrap_or(0);
-                                file_entry.insert("mod_time".to_string(), json!(mod_time));
+
+                        if !code.trim().is_empty() && !code.contains(char::REPLACEMENT_CHARACTER) {
+                            // ~~~ Filepath ~~~
+                            let file_path = if config.absolute_path {
+                                path.to_string_lossy().to_string()
+                            } else {
+                                relative_path.to_string_lossy().to_string()
+                            };
+
+                            // ~~~ File JSON Representation ~~~
+                            let mut file_entry = serde_json::Map::new();
+                            file_entry.insert("path".to_string(), json!(file_path));
+                            file_entry.insert(
+                                "extension".to_string(),
+                                json!(path.extension().and_then(|ext| ext.to_str()).unwrap_or("")),
+                            );
+                            file_entry.insert("code".to_string(), json!(code_block));
+
+                            // Store metadata
+                            let entry_meta = EntryMetadata::from(&metadata);
+                            file_entry
+                                .insert("metadata".to_string(), serde_json::to_value(entry_meta)?);
+
+                            // Add token count for the file only if token map is enabled
+                            if config.token_map_enabled {
+                                let token_count = count_tokens(&code, &config.encoding);
+                                file_entry.insert("token_count".to_string(), json!(token_count));
                             }
+
+                            // If date sorting is requested, record the file modification time.
+                            if let Some(method) = config.sort_method {
+                                if method == FileSortMethod::DateAsc
+                                    || method == FileSortMethod::DateDesc
+                                {
+                                    let mod_time = metadata
+                                        .modified()
+                                        .ok()
+                                        .and_then(|mtime| {
+                                            mtime
+                                                .duration_since(std::time::SystemTime::UNIX_EPOCH)
+                                                .ok()
+                                        })
+                                        .map(|d| d.as_secs())
+                                        .unwrap_or(0);
+                                    file_entry.insert("mod_time".to_string(), json!(mod_time));
+                                }
+                            }
+                            files.push(serde_json::Value::Object(file_entry));
+                            debug!(target: "included_files", "Included file: {}", file_path);
+                        } else {
+                            debug!("Excluded file (empty or invalid UTF-8): {}", path.display());
                         }
-                        files.push(serde_json::Value::Object(file_entry));
-                        debug!(target: "included_files", "Included file: {}", file_path);
                     } else {
-                        debug!("Excluded file (empty or invalid UTF-8): {}", path.display());
+                        debug!("Failed to read file: {}", path.display());
                     }
-                } else {
-                    debug!("Failed to read file: {}", path.display());
                 }
             }
         }

crates/code2prompt-core/src/template.rs

@@ -2,7 +2,9 @@
 //! It also includes functions for handling user-defined variables, copying the rendered output to the clipboard, and writing it to a file.
 use anyhow::{anyhow, Result};
 use colored::*;
-use handlebars::{no_escape, Handlebars};
+use handlebars::{
+    no_escape, Context, Handlebars, Helper, HelperResult, Output, RenderContext,
+};
 use regex::Regex;
 use std::io::Write;
 use std::str::FromStr;
@@ -21,6 +23,38 @@ pub fn handlebars_setup(template_str: &str, template_name: &str) -> Result<Handl
     let mut handlebars = Handlebars::new();
     handlebars.register_escape_fn(no_escape);
 
+    // Register custom helper for first N lines plus last line
+    handlebars.register_helper("firstLines", Box::new(|helper: &Helper, 
+                                                     _: &Handlebars,
+                                                     _: &Context,
+                                                     _: &mut RenderContext,
+                                                     out: &mut dyn Output| -> HelperResult {
+        // Get the text parameter
+        let text = helper.param(0).and_then(|v| v.value().as_str()).unwrap_or("");
+        // Get the count parameter
+        let count = helper.param(1).and_then(|v| v.value().as_u64()).unwrap_or(5) as usize;
+
+        let lines: Vec<&str> = text.lines().collect();
+
+        if lines.len() <= count + 1 {
+            // If there are fewer lines than count+1, just return the whole text
+            out.write(text)?;
+        } else {
+            // Take first N lines
+            let first_n_lines = lines.iter().take(count).cloned().collect::<Vec<&str>>();
+
+            // Get the last line
+            let last_line = lines.last().unwrap_or(&"");
+
+            // Combine first N lines and last line
+            let result = format!("{}\n{}", first_n_lines.join("\n"), last_line);
+
+            out.write(&result)?;
+        }
+
+        Ok(())
+    }));
+
     handlebars
         .register_template_string(template_name, template_str)
         .map_err(|e| anyhow!("Failed to register template: {}", e))?;