deepset-ai--haystack/test/components/preprocessors/test_python_code_splitter.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import textwrap

import pytest

from haystack import Document
from haystack.components.preprocessors import PythonCodeSplitter


@pytest.fixture
def simple_module_source():
    return textwrap.dedent(
        '''
        """Example module docstring."""
        import os
        import sys
        from math import sqrt


        def add(a, b):
            """Add two numbers."""
            return a + b


        def subtract(a, b):
            """Subtract two numbers."""
            return a - b
        '''
    ).lstrip()


@pytest.fixture
def class_source():
    return textwrap.dedent(
        '''
        """Geometry helpers."""
        from math import pi


        class Shape:
            """Base shape."""

            kind = "shape"

            def __init__(self, name: str) -> None:
                self.name = name

            def describe(self) -> str:
                return f"shape {self.name}"


        class Circle(Shape, metaclass=type):
            """A circle."""

            def __init__(self, r: float) -> None:
                super().__init__("circle")
                self.r = r

            @staticmethod
            def pi_value() -> float:
                return pi

            @classmethod
            def unit(cls) -> "Circle":
                return cls(1.0)

            def area(self) -> float:
                return pi * self.r * self.r
        '''
    ).lstrip()


@pytest.fixture
def oversized_function_source():
    """A function long enough to trigger the secondary line-based split."""
    body_lines = "\n".join(f"    x_{i} = {i}" for i in range(200))
    return f"def giant():\n{body_lines}\n    return x_0\n"


class TestInitValidation:
    def test_defaults(self):
        splitter = PythonCodeSplitter()
        assert splitter.min_effective_lines == 20
        assert splitter.max_effective_lines == 100
        assert splitter.expected_chars_per_line == 45
        assert splitter.oversized_factor == 3
        assert splitter.strip_docstrings is False
        assert splitter.preserve_class_definition is True
        assert splitter.secondary_split_overlap == 5
        assert splitter.secondary_split_length is None

    def test_custom_values(self):
        splitter = PythonCodeSplitter(
            min_effective_lines=2,
            max_effective_lines=10,
            expected_chars_per_line=80,
            oversized_factor=4,
            strip_docstrings=True,
            preserve_class_definition=False,
            secondary_split_overlap=2,
            secondary_split_length=15,
        )
        assert splitter.min_effective_lines == 2
        assert splitter.max_effective_lines == 10
        assert splitter.expected_chars_per_line == 80
        assert splitter.oversized_factor == 4
        assert splitter.strip_docstrings is True
        assert splitter.preserve_class_definition is False
        assert splitter.secondary_split_overlap == 2
        assert splitter.secondary_split_length == 15

    @pytest.mark.parametrize(
        "kwargs",
        [
            {"min_effective_lines": 0},
            {"min_effective_lines": -1},
            {"max_effective_lines": 0},
            {"max_effective_lines": -3},
            {"min_effective_lines": 10, "max_effective_lines": 5},
            {"expected_chars_per_line": 0},
            {"oversized_factor": 0},
            {"secondary_split_overlap": -1},
            {"secondary_split_length": -1},
        ],
    )
    def test_invalid_init_raises(self, kwargs):
        with pytest.raises(ValueError):
            PythonCodeSplitter(**kwargs)


class TestRunInputValidation:
    def test_none_content_raises_value_error(self):
        splitter = PythonCodeSplitter()
        doc = Document(content=None)
        with pytest.raises(ValueError):
            splitter.run(documents=[doc])

    def test_non_string_content_raises_type_error(self):
        splitter = PythonCodeSplitter()
        doc = Document(content="placeholder")
        doc.content = 12345  # type: ignore[assignment]
        with pytest.raises(TypeError):
            splitter.run(documents=[doc])

    def test_invalid_syntax_raises(self):
        splitter = PythonCodeSplitter()
        doc = Document(content="def broken(:\n    pass\n")
        with pytest.raises(SyntaxError):
            splitter.run(documents=[doc])

    def test_empty_documents_list(self):
        splitter = PythonCodeSplitter()
        result = splitter.run(documents=[])
        assert result == {"documents": []}


class TestBasicOutput:
    def test_returns_dict_with_documents(self, simple_module_source):
        splitter = PythonCodeSplitter()
        result = splitter.run(documents=[Document(content=simple_module_source)])
        assert isinstance(result, dict)
        assert "documents" in result
        assert isinstance(result["documents"], list)
        assert len(result["documents"]) >= 1
        for chunk in result["documents"]:
            assert isinstance(chunk, Document)
            assert isinstance(chunk.content, str)
            assert chunk.content

    def test_split_id_starts_at_zero_and_increments(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=class_source)])
        chunks = result["documents"]
        ids = [c.meta["split_id"] for c in chunks]
        assert ids == list(range(len(chunks)))

    def test_source_id_consistent_within_one_document(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=class_source)])
        chunks = result["documents"]
        source_ids = {c.meta["source_id"] for c in chunks}
        assert len(source_ids) == 1

    def test_source_id_differs_between_documents(self, simple_module_source, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
        docs = [Document(content=simple_module_source), Document(content=class_source)]
        result = splitter.run(documents=docs)
        source_ids = {c.meta["source_id"] for c in result["documents"]}
        assert len(source_ids) == 2

    def test_chunks_have_required_meta_fields(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
        result = splitter.run(documents=[Document(content=simple_module_source)])
        for chunk in result["documents"]:
            assert "source_id" in chunk.meta
            assert "split_id" in chunk.meta
            assert "start_line" in chunk.meta
            assert "end_line" in chunk.meta
            assert "unit_kinds" in chunk.meta
            assert isinstance(chunk.meta["unit_kinds"], list)
            assert chunk.meta["start_line"] >= 1
            assert chunk.meta["end_line"] >= chunk.meta["start_line"]

    def test_unit_kinds_lists_what_was_merged(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
        result = splitter.run(documents=[Document(content=simple_module_source)])
        all_kinds = set()
        for chunk in result["documents"]:
            for kind in chunk.meta["unit_kinds"]:
                all_kinds.add(kind)
        text = " ".join(all_kinds).lower()
        assert any("import" in t for t in all_kinds) or "import" in text
        assert any("func" in t or "method" in t for t in all_kinds) or any("func" in t for t in text.split())

    def test_multiple_documents_each_produces_chunks(self, simple_module_source, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
        docs = [
            Document(content=simple_module_source, meta={"file_name": "a.py"}),
            Document(content=class_source, meta={"file_name": "b.py"}),
        ]
        result = splitter.run(documents=docs)
        file_names = {c.meta.get("file_name") for c in result["documents"]}
        assert file_names == {"a.py", "b.py"}

    def test_split_id_resets_per_document(self, simple_module_source, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=10)
        docs = [
            Document(content=simple_module_source, meta={"file_name": "a.py"}),
            Document(content=class_source, meta={"file_name": "b.py"}),
        ]
        result = splitter.run(documents=docs)
        per_file_ids: dict[str, list[int]] = {"a.py": [], "b.py": []}
        for chunk in result["documents"]:
            per_file_ids[chunk.meta["file_name"]].append(chunk.meta["split_id"])
        for ids in per_file_ids.values():
            assert ids == sorted(ids)
            assert ids[0] == 0


class TestOrderingAndLineRanges:
    def test_chunks_are_in_source_order(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=class_source)])
        prev_start = 0
        for chunk in result["documents"]:
            assert chunk.meta["start_line"] >= prev_start
            assert chunk.meta["end_line"] >= chunk.meta["start_line"]
            prev_start = chunk.meta["start_line"]

    def test_chunks_dont_overlap_in_primary_split(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        chunks = splitter.run(documents=[Document(content=class_source)])["documents"]
        for prev, nxt in zip(chunks, chunks[1:], strict=False):
            assert nxt.meta["start_line"] > prev.meta["end_line"]

    def test_chunks_read_top_to_bottom(self, class_source):
        # With the default preserve_class_definition=True, chunks may have a class
        # signature prefixed but the source slice itself must still appear verbatim.
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=class_source)])
        source_lines = class_source.splitlines(keepends=True)
        for chunk in result["documents"]:
            expected = "".join(source_lines[chunk.meta["start_line"] - 1 : chunk.meta["end_line"]])
            assert expected in (chunk.content or "")

    def test_chunks_equal_source_slice_without_class_preservation(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, preserve_class_definition=False)
        result = splitter.run(documents=[Document(content=class_source)])
        source_lines = class_source.splitlines(keepends=True)
        for chunk in result["documents"]:
            expected = "".join(source_lines[chunk.meta["start_line"] - 1 : chunk.meta["end_line"]])
            assert (chunk.content or "").endswith(expected)


class TestFileNamePropagation:
    def test_file_name_propagated_to_all_chunks(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=simple_module_source, meta={"file_name": "sample.py"})])
        for chunk in result["documents"]:
            assert chunk.meta["file_name"] == "sample.py"

    def test_no_file_name_when_absent(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=simple_module_source)])
        for chunk in result["documents"]:
            assert "file_name" not in chunk.meta

    def test_other_meta_is_propagated(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5)
        result = splitter.run(
            documents=[Document(content=simple_module_source, meta={"file_name": "x.py", "project": "haystack"})]
        )
        for chunk in result["documents"]:
            assert chunk.meta["project"] == "haystack"


class TestDecorators:
    def test_decorators_metadata_present(self):
        source = textwrap.dedent(
            """
            class A:
                @staticmethod
                def s():
                    return 1

                @classmethod
                def c(cls):
                    return 2
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2)
        result = splitter.run(documents=[Document(content=source)])
        all_decorators = []
        for chunk in result["documents"]:
            all_decorators.extend(chunk.meta.get("decorators") or [])
        assert any("staticmethod" in d for d in all_decorators)
        assert any("classmethod" in d for d in all_decorators)

    def test_decorator_lines_included_in_chunk_content(self):
        source = textwrap.dedent(
            """
            class A:
                @staticmethod
                def s():
                    return 1
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2)
        result = splitter.run(documents=[Document(content=source)])
        chunks_with_s = [c for c in result["documents"] if "def s" in (c.content or "")]
        assert chunks_with_s
        for chunk in chunks_with_s:
            assert "@staticmethod" in (chunk.content or "")

    def test_decorators_deduped_in_chunk(self):
        source = textwrap.dedent(
            """
            class A:
                @staticmethod
                def one():
                    return 1

                @staticmethod
                def two():
                    return 2
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=20)
        result = splitter.run(documents=[Document(content=source)])
        for chunk in result["documents"]:
            decorators = chunk.meta.get("decorators") or []
            assert len(decorators) == len(set(decorators))

    def test_function_with_three_decorators_lists_all(self):
        source = textwrap.dedent(
            """
            def deco_a(fn):
                return fn


            def deco_b(fn):
                return fn


            def deco_c(fn):
                return fn


            @deco_a
            @deco_b
            @deco_c
            def triple():
                return 1
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
        result = splitter.run(documents=[Document(content=source)])
        chunks_with_triple = [c for c in result["documents"] if "def triple" in (c.content or "")]
        assert chunks_with_triple
        decorators = [d for c in chunks_with_triple for d in (c.meta.get("decorators") or [])]
        joined = " ".join(decorators)
        assert "deco_a" in joined
        assert "deco_b" in joined
        assert "deco_c" in joined

    def test_function_with_three_decorators_all_lines_in_content(self):
        source = textwrap.dedent(
            """
            @deco_a
            @deco_b
            @deco_c
            def triple():
                return 1
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
        result = splitter.run(documents=[Document(content=source)])

        chunks_with_triple = [c for c in result["documents"] if "def triple" in (c.content or "")]
        assert chunks_with_triple
        for chunk in chunks_with_triple:
            content = chunk.content or ""
            assert "@deco_a" in content
            assert "@deco_b" in content
            assert "@deco_c" in content


class TestIncludeClassesMeta:
    @pytest.mark.parametrize("class_name", ["Shape", "Circle"])
    def test_class_methods_carry_include_classes(self, class_source, class_name):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3)
        result = splitter.run(documents=[Document(content=class_source)])
        chunks = [c for c in result["documents"] if class_name in (c.meta.get("include_classes") or [])]
        assert chunks

    def test_include_classes_absent_when_no_class_involved(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
        result = splitter.run(documents=[Document(content=simple_module_source)])
        for chunk in result["documents"]:
            include_classes = chunk.meta.get("include_classes")
            assert not include_classes

    def test_include_classes_is_deduplicated(self):
        source = textwrap.dedent(
            """
            class A:
                def one(self):
                    return 1

                def two(self):
                    return 2

                def three(self):
                    return 3
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=200)
        result = splitter.run(documents=[Document(content=source)])
        for chunk in result["documents"]:
            include_classes = chunk.meta.get("include_classes") or []
            assert len(include_classes) == len(set(include_classes))

    def test_include_classes_preserves_source_order(self):
        source = textwrap.dedent(
            """
            class First:
                def f(self):
                    return 1


            class Second:
                def g(self):
                    return 2


            class Third:
                def h(self):
                    return 3
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=500)
        result = splitter.run(documents=[Document(content=source)])

        for chunk in result["documents"]:
            include_classes = chunk.meta.get("include_classes") or []
            if len(include_classes) >= 2:
                expected_order = [c for c in ["First", "Second", "Third"] if c in include_classes]
                assert include_classes == expected_order


class TestPreserveClassDefinition:
    @pytest.fixture
    def multi_method_class_source(self):
        return textwrap.dedent(
            '''
            class Greeter:
                """A friendly greeter."""

                kind = "greeter"

                def __init__(self, name: str) -> None:
                    self.name = name

                def hello(self) -> str:
                    return f"hello {self.name}"

                def bye(self) -> str:
                    return f"bye {self.name}"

                def shout(self) -> str:
                    return f"HELLO {self.name.upper()}"

                def whisper(self) -> str:
                    return f"hello {self.name}..."
            '''
        ).lstrip()

    def test_default_preserve_class_definition_is_true(self):
        splitter = PythonCodeSplitter()
        assert splitter.preserve_class_definition is True

    def test_class_signature_prepended_to_later_chunks(self, multi_method_class_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
        chunks = splitter.run(documents=[Document(content=multi_method_class_source)])["documents"]
        assert len(chunks) >= 2
        for chunk in chunks:
            if "Greeter" in (chunk.meta.get("include_classes") or []):
                assert "class Greeter" in (chunk.content or "")

    def test_disabled_does_not_prepend_class_signature(self, multi_method_class_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=False)
        chunks = splitter.run(documents=[Document(content=multi_method_class_source)])["documents"]
        chunks_with_header = [c for c in chunks if "class Greeter" in (c.content or "")]
        assert len(chunks_with_header) == 1

    def test_preserve_does_not_duplicate_header_in_original_chunk(self, multi_method_class_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
        result = splitter.run(documents=[Document(content=multi_method_class_source)])
        for chunk in result["documents"]:
            assert (chunk.content or "").count("class Greeter") <= 1

    def test_preserve_keeps_inheritance_and_metaclass(self):
        source = textwrap.dedent(
            '''
            class Base:
                pass


            class Meta(type):
                pass


            class Child(Base, metaclass=Meta):
                """A child class."""

                def one(self):
                    return 1

                def two(self):
                    return 2

                def three(self):
                    return 3

                def four(self):
                    return 4
            '''
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
        result = splitter.run(documents=[Document(content=source)])
        child_chunks = [c for c in result["documents"] if "Child" in (c.meta.get("include_classes") or [])]
        assert len(child_chunks) >= 2
        for chunk in child_chunks:
            assert "class Child(Base, metaclass=Meta):" in (chunk.content or "")

    def test_preserve_keeps_decorators_on_class(self):
        source = textwrap.dedent(
            """
            def reg(cls):
                return cls


            @reg
            class Decorated:
                def one(self):
                    return 1

                def two(self):
                    return 2

                def three(self):
                    return 3

                def four(self):
                    return 4
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
        result = splitter.run(documents=[Document(content=source)])
        decorated_chunks = [c for c in result["documents"] if "Decorated" in (c.meta.get("include_classes") or [])]
        assert len(decorated_chunks) >= 2
        for chunk in decorated_chunks:
            content = chunk.content or ""
            assert "class Decorated" in content
            assert "@reg" in content

    def test_preserve_handles_multiple_classes(self):
        source = textwrap.dedent(
            """
            class A:
                def a1(self):
                    return 1

                def a2(self):
                    return 2

                def a3(self):
                    return 3


            class B:
                def b1(self):
                    return 1

                def b2(self):
                    return 2

                def b3(self):
                    return 3
            """
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=3, preserve_class_definition=True)
        result = splitter.run(documents=[Document(content=source)])
        for chunk in result["documents"]:
            content = chunk.content or ""
            for cls in chunk.meta.get("include_classes") or []:
                assert f"class {cls}" in content


class TestDocstringStripping:
    def test_default_keeps_docstrings_in_content(self, class_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=8)
        result = splitter.run(documents=[Document(content=class_source)])
        contents = "\n".join(c.content or "" for c in result["documents"])
        assert "A circle." in contents
        for chunk in result["documents"]:
            assert not chunk.meta.get("docstrings")

    def test_strip_docstrings_moves_them_to_meta(self):
        source = textwrap.dedent(
            '''
            def foo():
                """Foo docstring."""
                return 1


            def bar():
                """Bar docstring."""
                return 2
            '''
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=10, strip_docstrings=True)
        result = splitter.run(documents=[Document(content=source)])
        all_docstrings = [d for c in result["documents"] for d in (c.meta.get("docstrings") or [])]
        joined_docstrings = " | ".join(all_docstrings)
        assert "Foo docstring." in joined_docstrings
        assert "Bar docstring." in joined_docstrings
        joined_content = "\n".join(c.content or "" for c in result["documents"])
        assert "Foo docstring." not in joined_content
        assert "Bar docstring." not in joined_content

    def test_strip_docstrings_preserves_module_docstring(self):
        source = textwrap.dedent(
            '''
            """Module-level docstring."""

            def foo():
                """Inner."""
                return 1
            '''
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=20, strip_docstrings=True)
        result = splitter.run(documents=[Document(content=source)])
        joined = "\n".join(c.content or "" for c in result["documents"])
        assert "Module-level docstring." in joined

    def test_strip_class_header_docstring_moves_to_meta(self):
        source = textwrap.dedent(
            '''
            class MyClass:
                """Class-level docstring."""

                class_var = 42

                def method(self):
                    return self.class_var
            '''
        ).lstrip()
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=10, strip_docstrings=True)
        result = splitter.run(documents=[Document(content=source)])
        header_chunks = [c for c in result["documents"] if "class_header" in c.meta.get("unit_kinds", [])]
        assert header_chunks
        header = header_chunks[0]
        assert "Class-level docstring." not in (header.content or "")
        assert "Class-level docstring." in " | ".join(header.meta.get("docstrings") or [])


class TestTopLevelStatements:
    @pytest.fixture
    def rich_module_source(self):
        return textwrap.dedent(
            '''
            """Utility helpers for the pipeline."""
            import os
            import sys
            from pathlib import Path

            MAX_RETRIES = 3
            DEFAULT_TIMEOUT = 30.0
            LOG_PREFIX = "app"


            def process(data):
                """Process data."""
                result = data.strip()
                return result


            def validate(value):
                """Validate a value."""
                if value is None:
                    raise ValueError("value cannot be None")
                return True


            class Manager:
                """Resource manager."""

                def __init__(self):
                    self.items = []

                def add(self, item):
                    self.items.append(item)

                def remove(self, item):
                    self.items.remove(item)


            if __name__ == "__main__":
                mgr = Manager()
                mgr.add(process("test"))
            '''
        ).lstrip()

    @pytest.mark.parametrize("expected_kind", ["statement", "imports", "module_docstring"])
    def test_unit_kind_present(self, rich_module_source, expected_kind):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
        result = splitter.run(documents=[Document(content=rich_module_source)])
        all_kinds = [k for c in result["documents"] for k in c.meta.get("unit_kinds", [])]
        assert expected_kind in all_kinds

    def test_first_chunk_contains_preamble_statements(self, rich_module_source):
        # max_effective_lines=6 matches the preamble (module docstring + 3 imports + 3
        # assignments) so the greedy merger flushes before the first function.
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=6)
        result = splitter.run(documents=[Document(content=rich_module_source)])
        assert len(result["documents"]) >= 2
        first = result["documents"][0]
        content = first.content or ""
        assert '"""Utility helpers for the pipeline."""' in content
        assert "import os" in content
        assert "import sys" in content
        assert "from pathlib import Path" in content
        assert "MAX_RETRIES = 3" in content
        assert "DEFAULT_TIMEOUT = 30.0" in content
        assert 'LOG_PREFIX = "app"' in content
        assert "def process" not in content
        assert "def validate" not in content
        assert "class Manager" not in content

    def test_if_main_produces_statement_unit(self, rich_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
        result = splitter.run(documents=[Document(content=rich_module_source)])
        all_kinds = [k for c in result["documents"] for k in c.meta.get("unit_kinds", [])]
        assert "statement" in all_kinds
        joined = "\n".join(c.content or "" for c in result["documents"])
        assert 'if __name__ == "__main__"' in joined

    def test_all_imports_appear_in_output(self, rich_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
        result = splitter.run(documents=[Document(content=rich_module_source)])
        joined = "\n".join(c.content or "" for c in result["documents"])
        assert "import os" in joined
        assert "import sys" in joined
        assert "from pathlib import Path" in joined

    def test_module_docstring_text_preserved(self, rich_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=50)
        result = splitter.run(documents=[Document(content=rich_module_source)])
        joined = "\n".join(c.content or "" for c in result["documents"])
        assert "Utility helpers for the pipeline." in joined


class TestOversizedFallback:
    def test_warns_on_oversized_function(self, oversized_function_source, caplog):
        import logging

        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
        with caplog.at_level(logging.WARNING):
            result = splitter.run(documents=[Document(content=oversized_function_source)])
        text = caplog.text.lower()
        assert "oversiz" in text or "secondary" in text
        assert len(result["documents"]) >= 2

    def test_oversized_chunks_marked_with_secondary_split(self, oversized_function_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
        result = splitter.run(documents=[Document(content=oversized_function_source)])
        assert [c for c in result["documents"] if c.meta.get("secondary_split")]

    def test_oversized_chunks_belong_to_same_function(self, oversized_function_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=5, oversized_factor=3)
        result = splitter.run(documents=[Document(content=oversized_function_source)])
        secondary = [c for c in result["documents"] if c.meta.get("secondary_split")]
        for chunk in secondary:
            assert chunk.meta["start_line"] >= 1
            assert chunk.meta["end_line"] <= oversized_function_source.count("\n") + 1

    def test_small_function_does_not_trigger_secondary(self, simple_module_source):
        splitter = PythonCodeSplitter(min_effective_lines=2, max_effective_lines=50, oversized_factor=3)
        result = splitter.run(documents=[Document(content=simple_module_source)])
        for chunk in result["documents"]:
            assert not chunk.meta.get("secondary_split")

    def test_long_lines_count_as_more_effective_lines(self):
        long_line_source = (
            textwrap.dedent(
                """
            def short():
                return 1


            def longer():
                return "{padding}"
            """
            )
            .lstrip()
            .format(padding="x" * 500)
        )
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=2, expected_chars_per_line=10)
        result = splitter.run(documents=[Document(content=long_line_source)])
        chunks_with_short = [c for c in result["documents"] if "def short" in (c.content or "")]
        chunks_with_long = [c for c in result["documents"] if "def longer" in (c.content or "")]
        assert chunks_with_short and chunks_with_long
        assert chunks_with_short[0] is not chunks_with_long[0]


class TestEdgeCases:
    def test_module_with_only_docstring(self):
        source = '"""Just a docstring."""\n'
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=source)])
        assert len(result["documents"]) == 1
        assert "Just a docstring." in (result["documents"][0].content or "")

    def test_module_with_only_imports(self):
        source = "import os\nimport sys\nfrom math import pi\n"
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=source)])
        joined = "\n".join(c.content or "" for c in result["documents"])
        assert "import os" in joined
        assert "import sys" in joined
        assert "from math import pi" in joined

    def test_module_with_only_one_function(self):
        source = "def f():\n    return 1\n"
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
        result = splitter.run(documents=[Document(content=source)])
        assert len(result["documents"]) == 1
        assert "def f" in (result["documents"][0].content or "")

    def test_empty_string_source(self):
        # Empty source is valid Python; the splitter must not raise.
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
        result = splitter.run(documents=[Document(content="")])
        assert isinstance(result["documents"], list)

    def test_invalid_syntax_raises_syntax_error(self):
        splitter = PythonCodeSplitter(min_effective_lines=1, max_effective_lines=5)
        doc = Document(content="class Broken(\n    pass\n")
        with pytest.raises(SyntaxError):
            splitter.run(documents=[doc])