diff --git a/spacy/tests/doc/test_span_group.py b/spacy/tests/doc/test_span_group.py index 8c70a83e1..818569c64 100644 --- a/spacy/tests/doc/test_span_group.py +++ b/spacy/tests/doc/test_span_group.py @@ -1,7 +1,10 @@ +from typing import List + import pytest from random import Random from spacy.matcher import Matcher -from spacy.tokens import Span, SpanGroup +from spacy.tokens import Span, SpanGroup, Doc +from spacy.util import filter_spans @pytest.fixture @@ -240,3 +243,13 @@ def test_span_group_extend(doc): def test_span_group_dealloc(span_group): with pytest.raises(AttributeError): print(span_group.doc) + + +@pytest.mark.issue(11975) +def test_span_group_typing(doc: Doc): + """Tests whether typing of `SpanGroup` as `Iterable[Span]`-like object is accepted by mypy.""" + span_group: SpanGroup = doc.spans["SPANS"] + spans: List[Span] = list(span_group) + for i, span in enumerate(span_group): + assert span == span_group[i] == spans[i] + filter_spans(span_group) diff --git a/spacy/tokens/span.pyi b/spacy/tokens/span.pyi index 0a6f306a6..9986a90e6 100644 --- a/spacy/tokens/span.pyi +++ b/spacy/tokens/span.pyi @@ -95,8 +95,8 @@ class Span: self, start_idx: int, end_idx: int, - label: int = ..., - kb_id: int = ..., + label: Union[int, str] = ..., + kb_id: Union[int, str] = ..., vector: Optional[Floats1d] = ..., ) -> Span: ... @property diff --git a/spacy/tokens/span_group.pyi b/spacy/tokens/span_group.pyi index 21cd124ab..0b4aa83aa 100644 --- a/spacy/tokens/span_group.pyi +++ b/spacy/tokens/span_group.pyi @@ -18,6 +18,7 @@ class SpanGroup: def doc(self) -> Doc: ... @property def has_overlap(self) -> bool: ... + def __iter__(self): ... def __len__(self) -> int: ... def append(self, span: Span) -> None: ... def extend(self, spans: Iterable[Span]) -> None: ... diff --git a/spacy/tokens/span_group.pyx b/spacy/tokens/span_group.pyx index 1aa3c0bc8..608dda283 100644 --- a/spacy/tokens/span_group.pyx +++ b/spacy/tokens/span_group.pyx @@ -158,6 +158,16 @@ cdef class SpanGroup: return self._concat(other) return NotImplemented + def __iter__(self): + """ + Iterate over the spans in this SpanGroup. + YIELDS (Span): A span in this SpanGroup. + + DOCS: https://spacy.io/api/spangroup#iter + """ + for i in range(self.c.size()): + yield self[i] + def append(self, Span span): """Add a span to the group. The span must refer to the same Doc object as the span group. diff --git a/website/docs/api/spangroup.md b/website/docs/api/spangroup.md index 2d1cf73c4..bd9659acb 100644 --- a/website/docs/api/spangroup.md +++ b/website/docs/api/spangroup.md @@ -202,6 +202,23 @@ already present in the current span group. | `other` | The span group or spans to append. ~~Union[SpanGroup, Iterable[Span]]~~ | | **RETURNS** | The span group. ~~SpanGroup~~ | +## SpanGroup.\_\_iter\_\_ {#iter tag="method" new="3.5"} + +Iterate over the spans in this span group. + +> #### Example +> +> ```python +> doc = nlp("Their goi ng home") +> doc.spans["errors"] = [doc[0:1], doc[1:3]] +> for error_span in doc.spans["errors"]: +> print(error_span) +> ``` + +| Name | Description | +| ---------- | ----------------------------------- | +| **YIELDS** | A span in this span group. ~~Span~~ | + ## SpanGroup.append {#append tag="method"} Add a [`Span`](/api/span) object to the group. The span must refer to the same