nltk--nltk/nltk/test/unit/test_projectivedependencyparser.py

"""Regression tests for the quadratic-memory / cubic-time DoS in
``nltk.parse.projectivedependencyparser`` (CWE-407).

Both ``ProjectiveDependencyParser`` and ``ProbabilisticProjectiveDependencyParser``
eagerly build a dense ``(n+1)x(n+1)`` chart (O(n^2) memory) and run an
unconditional triple loop over it (O(n^3) time), regardless of the grammar -- so
a short list of a few hundred/thousand tokens pins a CPU core and exhausts
memory. ``parse`` now refuses with a ``ValueError`` once the token count exceeds
the configurable ``MAX_TOKENS``, before any chart is allocated; ordinary
sentence-length parses are unchanged.
"""

import pytest

from nltk.grammar import DependencyGrammar
from nltk.parse.projectivedependencyparser import (
    ProbabilisticProjectiveDependencyParser,
    ProjectiveDependencyParser,
)

_GRAMMAR = DependencyGrammar.fromstring(
    """
    'fell' -> 'price' | 'stock'
    'price' -> 'of' | 'the'
    'of' -> 'stock'
    'stock' -> 'the'
    """
)
_SENTENCE = ["the", "price", "of", "the", "stock", "fell"]


def test_default_max_tokens_is_positive_int():
    assert isinstance(ProjectiveDependencyParser.MAX_TOKENS, int)
    assert ProjectiveDependencyParser.MAX_TOKENS > 0
    assert isinstance(ProbabilisticProjectiveDependencyParser.MAX_TOKENS, int)
    assert ProbabilisticProjectiveDependencyParser.MAX_TOKENS > 0


def test_projective_parse_preserves_results():
    """An ordinary sentence parses to the same trees as before."""
    parser = ProjectiveDependencyParser(_GRAMMAR)
    trees = list(parser.parse(_SENTENCE))
    assert len(trees) == 3
    assert "(fell (price the of the) stock)" in [str(t) for t in trees]


def test_projective_parse_accepts_generator_input():
    """A non-subscriptable iterable (generator) parses like a list.

    parse() materializes the input with list(...) but later indexes tokens[i]
    when building output; the local name is rebound to the materialized list so
    an iterator/generator input does not raise TypeError.
    """
    parser = ProjectiveDependencyParser(_GRAMMAR)
    trees = list(parser.parse(iter(_SENTENCE)))
    assert len(trees) == 3
    assert "(fell (price the of the) stock)" in [str(t) for t in trees]


def test_projective_parse_rejects_over_cap():
    """Over the default cap, parse refuses instead of building the chart."""
    parser = ProjectiveDependencyParser(DependencyGrammar.fromstring("'x' -> 'y'"))
    n = ProjectiveDependencyParser.MAX_TOKENS + 1
    with pytest.raises(ValueError, match="MAX_TOKENS"):
        list(parser.parse(["a"] * n))


def test_projective_parse_respects_lowered_cap():
    """MAX_TOKENS is tunable per instance; at/under the cap still parses."""
    parser = ProjectiveDependencyParser(DependencyGrammar.fromstring("'x' -> 'y'"))
    parser.MAX_TOKENS = 3
    with pytest.raises(ValueError, match="MAX_TOKENS"):
        list(parser.parse(["a", "a", "a", "a"]))
    # exactly at the cap is allowed (no match -> empty result, but no error).
    assert list(parser.parse(["a", "a", "a"])) == []


def test_probabilistic_parse_rejects_over_cap():
    """The probabilistic decoder enforces the cap before touching the grammar."""
    parser = ProbabilisticProjectiveDependencyParser()
    parser.MAX_TOKENS = 4
    with pytest.raises(ValueError, match="MAX_TOKENS"):
        list(parser.parse(["a"] * 5))