Add --line-numbers flag (simonw#38)

* add --line-numbers flag
* -n/--line-numbers in README, refs simonw#38

Co-authored-by: Simon Willison <swillison@gmail.com>

Author: danclaytondev

README.md

@@ -64,6 +64,22 @@ This will output the contents of every file, with each file preceded by its rela
   files-to-prompt path/to/directory -o output.txt
   ```
 
+- `-n/--line-numbers`: Include line numbers in the output.
+
+  ```bash
+  files-to-prompt path/to/directory -n
+  ```
+  Example output:
+  ```
+  files_to_prompt/cli.py
+  ---
+    1  import os
+    2  from fnmatch import fnmatch
+    3
+    4  import click
+    ...
+  ```
+
 ### Example
 
 Suppose you have a directory structure like this:

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>")

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