From d77edc4350e6d9b55c857f71c5a49e7f6f910aca Mon Sep 17 00:00:00 2001
From: Dan Clayton <daniel@cl8on.co.uk>
Date: Tue, 21 Jan 2025 16:47:29 +0000
Subject: [PATCH] add --line-numbers flag

---
 files_to_prompt/cli.py        | 39 ++++++++++++++++++++++++++++-------
 tests/test_files_to_prompt.py | 28 +++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/files_to_prompt/cli.py b/files_to_prompt/cli.py
index 4fb358b..5b5423b 100644
--- a/files_to_prompt/cli.py
+++ b/files_to_prompt/cli.py
@@ -25,26 +25,39 @@ def read_gitignore(path):
     return []
 
 
-def print_path(writer, path, content, xml):
+def add_line_numbers(content):
+    lines = content.splitlines()
+
+    padding = len(str(len(lines)))
+
+    numbered_lines = [f"{i+1:{padding}}  {line}" for i, line in enumerate(lines)]
+    return "\n".join(numbered_lines)
+
+
+def print_path(writer, path, content, xml, line_numbers):
     if xml:
-        print_as_xml(writer, path, content)
+        print_as_xml(writer, path, content, line_numbers)
     else:
-        print_default(writer, path, content)
+        print_default(writer, path, content, line_numbers)
 
 
-def print_default(writer, path, content):
+def print_default(writer, path, content, line_numbers):
     writer(path)
     writer("---")
+    if line_numbers:
+        content = add_line_numbers(content)
     writer(content)
     writer("")
     writer("---")
 
 
-def print_as_xml(writer, path, content):
+def print_as_xml(writer, path, content, line_numbers):
     global global_index
     writer(f'<document index="{global_index}">')
     writer(f"<source>{path}</source>")
     writer("<document_content>")
+    if line_numbers:
+        content = add_line_numbers(content)
     writer(content)
     writer("</document_content>")
     writer("</document>")
@@ -60,11 +73,12 @@ def process_path(
     ignore_patterns,
     writer,
     claude_xml,
+    line_numbers=False,
 ):
     if os.path.isfile(path):
         try:
             with open(path, "r") as f:
-                print_path(writer, path, f.read(), claude_xml)
+                print_path(writer, path, f.read(), claude_xml, line_numbers)
         except UnicodeDecodeError:
             warning_message = f"Warning: Skipping file {path} due to UnicodeDecodeError"
             click.echo(click.style(warning_message, fg="red"), err=True)
@@ -101,7 +115,9 @@ def process_path(
                 file_path = os.path.join(root, file)
                 try:
                     with open(file_path, "r") as f:
-                        print_path(writer, file_path, f.read(), claude_xml)
+                        print_path(
+                            writer, file_path, f.read(), claude_xml, line_numbers
+                        )
                 except UnicodeDecodeError:
                     warning_message = (
                         f"Warning: Skipping file {file_path} due to UnicodeDecodeError"
@@ -143,6 +159,13 @@ def process_path(
     is_flag=True,
     help="Output in XML-ish format suitable for Claude's long context window.",
 )
+@click.option(
+    "line_numbers",
+    "-n",
+    "--line-numbers",
+    is_flag=True,
+    help="Add line numbers to the output",
+)
 @click.version_option()
 def cli(
     paths,
@@ -152,6 +175,7 @@ def cli(
     ignore_patterns,
     output_file,
     claude_xml,
+    line_numbers,
 ):
     """
     Takes one or more paths to files or directories and outputs every file,
@@ -204,6 +228,7 @@ def cli(
             ignore_patterns,
             writer,
             claude_xml,
+            line_numbers,
         )
     if claude_xml:
         writer("</documents>")
diff --git a/tests/test_files_to_prompt.py b/tests/test_files_to_prompt.py
index 207a515..24d7234 100644
--- a/tests/test_files_to_prompt.py
+++ b/tests/test_files_to_prompt.py
@@ -279,3 +279,31 @@ def test_output_option(tmpdir, arg):
 ---
 """
         assert expected.strip() == actual.strip()
+
+
+def test_line_numbers(tmpdir):
+    runner = CliRunner()
+    with tmpdir.as_cwd():
+        os.makedirs("test_dir")
+        test_content = "First line\nSecond line\nThird line\nFourth line\n"
+        with open("test_dir/multiline.txt", "w") as f:
+            f.write(test_content)
+
+        result = runner.invoke(cli, ["test_dir"])
+        assert result.exit_code == 0
+        assert "1  First line" not in result.output
+        assert test_content in result.output
+
+        result = runner.invoke(cli, ["test_dir", "-n"])
+        assert result.exit_code == 0
+        assert "1  First line" in result.output
+        assert "2  Second line" in result.output
+        assert "3  Third line" in result.output
+        assert "4  Fourth line" in result.output
+
+        result = runner.invoke(cli, ["test_dir", "--line-numbers"])
+        assert result.exit_code == 0
+        assert "1  First line" in result.output
+        assert "2  Second line" in result.output
+        assert "3  Third line" in result.output
+        assert "4  Fourth line" in result.output