Merge pull request #99 from mufeedvh/python-sdk

Code2prompt should have a SDK with uniform API

Author: ODAncona

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "code2prompt_core"
-version = "3.1.0"
+version = "3.2.0"
 authors = [
     "Mufeed VH <[email protected]>",
     "Olivier D'Ancona <[email protected]>",

crates/code2prompt-core/Cargo.toml

@@ -6,6 +6,7 @@ use crate::template::OutputFormat;
 use crate::tokenizer::TokenizerType;
 use crate::{sort::FileSortMethod, tokenizer::TokenFormat};
 use derive_builder::Builder;
+use std::collections::HashMap;
 use std::path::PathBuf;
 
 /// A stateless configuration object describing all the preferences and filters
@@ -33,9 +34,9 @@ pub struct Code2PromptConfig {
     #[builder(default)]
     pub line_numbers: bool,
 
-    /// If true, paths in the output will be relative instead of absolute.
+    /// If true, paths in the output will be absolute instead of relative.
     #[builder(default)]
-    pub relative_paths: bool,
+    pub absolute_path: bool,
 
     /// If true, code2prompt will generate a full directory tree, ignoring include/exclude rules.
     #[builder(default)]
@@ -96,6 +97,10 @@ pub struct Code2PromptConfig {
     /// The template string itself.
     #[builder(default)]
     pub template_str: String,
+
+    /// Extra template data
+    #[builder(default)]
+    pub user_variables: HashMap<String, String>,
 }
 
 impl Code2PromptConfig {

crates/code2prompt-core/src/configuration.rs

@@ -2,7 +2,7 @@ Project Path: {{ absolute_code_path }}
 
 Source Tree:
 
-```
+```txt
 {{ source_tree }}
 ```
 

crates/default_template_md.hbs renamed to crates/code2prompt-core/src/default_template_md.hbs

@@ -20,7 +20,7 @@ use termtree::Tree;
 /// * `exclude` - The patterns of files to exclude.
 /// * `include_priority` - Whether to give priority to include patterns.
 /// * `line_number` - Whether to add line numbers to the code.
-/// * `relative_paths` - Whether to use relative paths.
+/// * `absolute_path` - Whether to use absolute paths.
 ///
 /// # Returns
 ///
@@ -99,10 +99,10 @@ pub fn traverse_directory(config: &Code2PromptConfig) -> Result<(String, Vec<ser
 
                     if !code.trim().is_empty() && !code.contains(char::REPLACEMENT_CHARACTER) {
                         // ~~~ Filepath ~~~
-                        let file_path = if config.relative_paths {
-                            format!("{}/{}", parent_directory, relative_path.display())
-                        } else {
+                        let file_path = if config.absolute_path {
                             path.display().to_string()
+                        } else {
+                            format!("{}/{}", parent_directory, relative_path.display())
                         };
 
                         // ~~~ File JSON Representation ~~~

crates/default_template_xml.hbs renamed to crates/code2prompt-core/src/default_template_xml.hbs

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 use crate::configuration::Code2PromptConfig;
 use crate::git::{get_git_diff, get_git_diff_between_branches, get_git_log};
 use crate::path::{label, traverse_directory};
-use crate::template::{handlebars_setup, render_template};
+use crate::template::{handlebars_setup, render_template, OutputFormat};
 use crate::tokenizer::{count_tokens, TokenizerType};
 
 /// Represents a live session that holds stateful data about the user's codebase,
@@ -101,24 +101,50 @@ impl Code2PromptSession {
 
     /// Constructs a JSON object that merges the session data and your config’s path label.
     pub fn build_template_data(&self) -> serde_json::Value {
-        serde_json::json!({
+        let mut data = serde_json::json!({
             "absolute_code_path": label(&self.config.path),
             "source_tree": self.data.source_tree,
             "files": self.data.files,
             "git_diff": self.data.git_diff,
             "git_diff_branch": self.data.git_diff_branch,
             "git_log_branch": self.data.git_log_branch
-        })
+        });
+
+        // Add user-defined variables to the template data
+        if self.config.user_variables.len() > 0 {
+            if let Some(obj) = data.as_object_mut() {
+                for (key, value) in &self.config.user_variables {
+                    obj.insert(key.clone(), serde_json::Value::String(value.clone()));
+                }
+            }
+        }
+
+        data
     }
 
     /// Renders the final prompt given a template-data JSON object. Returns both
     /// the rendered prompt and the token count information. The session
     /// does not do any printing or user prompting — that’s up to the caller.
     pub fn render_prompt(&self, template_data: &serde_json::Value) -> Result<RenderedPrompt> {
+        // ~~~ Template selection ~~~
+        let mut template_str = self.config.template_str.clone();
+        let mut template_name = self.config.template_name.clone();
+        if self.config.template_str.is_empty() {
+            template_str = match self.config.output_format {
+                OutputFormat::Markdown | OutputFormat::Json => {
+                    include_str!("./default_template_md.hbs").to_string()
+                }
+                OutputFormat::Xml => include_str!("./default_template_xml.hbs").to_string(),
+            };
+            template_name = match self.config.output_format {
+                OutputFormat::Markdown | OutputFormat::Json => "markdown".to_string(),
+                OutputFormat::Xml => "xml".to_string(),
+            };
+        }
+
         // ~~~ Rendering ~~~
-        let handlebars = handlebars_setup(&self.config.template_str, &self.config.template_name)?;
-        let rendered_prompt =
-            render_template(&handlebars, &self.config.template_name, template_data)?;
+        let handlebars = handlebars_setup(&template_str, &template_name)?;
+        let rendered_prompt = render_template(&handlebars, &template_name, template_data)?;
 
         // ~~~ Informations ~~~
         let tokenizer_type: TokenizerType = self.config.encoding;
@@ -153,9 +179,30 @@ impl Code2PromptSession {
 
     pub fn generate_prompt(&mut self) -> Result<RenderedPrompt> {
         self.load_codebase()?;
-        self.load_git_diff()?;
-        self.load_git_diff_between_branches()?;
-        self.load_git_log_between_branches()?;
+
+        // ~~~~ Load Git info ~~~
+        if self.config.diff_enabled {
+            match self.load_git_diff() {
+                Ok(_) => {}
+                Err(e) => log::warn!("Git diff could not be loaded: {}", e),
+            }
+        }
+
+        // ~~~ Load Git info between branches ~~~
+        if self.config.diff_branches.is_some() {
+            match self.load_git_diff_between_branches() {
+                Ok(_) => {}
+                Err(e) => log::warn!("Git branch diff could not be loaded: {}", e),
+            }
+        }
+
+        // ~~~ Load Git log between branches ~~~
+        if self.config.log_branches.is_some() {
+            match self.load_git_log_between_branches() {
+                Ok(_) => {}
+                Err(e) => log::warn!("Git branch log could not be loaded: {}", e),
+            }
+        }
         let template_data = self.build_template_data();
         let rendered = self.render_prompt(&template_data)?;
         Ok(rendered)

crates/code2prompt-core/src/path.rs

@@ -1,6 +1,6 @@
 [package]
 name = "code2prompt-python"
-version = "3.1.0"
+version = "3.2.0"
 edition = "2021"
 
 [lib]

crates/code2prompt-core/src/session.rs

@@ -1,17 +1,12 @@
 [project]
 name = "code2prompt_rs"
-version = "3.0.0"
+version = "3.2.0"
 description = "Python bindings for code2prompt"
 authors = [
-    { name = "Mufeed VH", email = "[email protected]" },
     { name = "Olivier D'Ancona", email = "[email protected]" },
+    { name = "Mufeed VH", email = "[email protected]" },
 ]
-dependencies = [
-    "maturin>=1.8.2",
-    "pytest>=8.3.5",
-    "pip>=25.0.1",
-    "patchelf>=0.17.2.1",
-]
+dependencies = ["maturin>=1.8.2", "pip>=25.0.1", "patchelf>=0.17.2.1"]
 requires-python = ">= 3.11"
 classifiers = [
     "Development Status :: 5 - Production/Stable",
@@ -38,7 +33,7 @@ features = ["pyo3/extension-module"]
 
 [tool.rye]
 managed = true
-dev-dependencies = ["maturin>=1.8.2"]
+dev-dependencies = ["maturin>=1.8.2", "pytest>=8.3.5"]
 
 [tool.rye.scripts]
 build = "maturin develop"

crates/code2prompt-python/Cargo.toml

@@ -8,7 +8,7 @@
 - Count tokens for different models
 """
 
-from .code2prompt import CodePrompt
+# Import the Python wrapper class from the renamed file
+from .code2prompt_rs import Code2Prompt
 
-__version__ = "2.0.0"
-__all__ = ["CodePrompt"]
+__all__ = ['Code2Prompt']