rfctr: improve typing local to BlockItemContainer

Author: scanny

features/steps/block.py

@@ -1,6 +1,7 @@
 """Step implementations for block content containers."""
 
 from behave import given, then, when
+from behave.runner import Context
 
 from docx import Document
 from docx.table import Table
@@ -11,12 +12,12 @@
 
 
 @given("a document containing a table")
-def given_a_document_containing_a_table(context):
+def given_a_document_containing_a_table(context: Context):
  context.document = Document(test_docx("blk-containing-table"))
 
 
 @given("a paragraph")
-def given_a_paragraph(context):
+def given_a_paragraph(context: Context):
  context.document = Document()
  context.paragraph = context.document.add_paragraph()
 
@@ -25,13 +26,13 @@ def given_a_paragraph(context):
 
 
 @when("I add a paragraph")
-def when_add_paragraph(context):
+def when_add_paragraph(context: Context):
  document = context.document
  context.p = document.add_paragraph()
 
 
 @when("I add a table")
-def when_add_table(context):
+def when_add_table(context: Context):
  rows, cols = 2, 2
  context.document.add_table(rows, cols)
 
@@ -40,12 +41,12 @@ def when_add_table(context):
 
 
 @then("I can access the table")
-def then_can_access_table(context):
+def then_can_access_table(context: Context):
  table = context.document.tables[-1]
  assert isinstance(table, Table)
 
 
 @then("the new table appears in the document")
-def then_new_table_appears_in_document(context):
+def then_new_table_appears_in_document(context: Context):
  table = context.document.tables[-1]
  assert isinstance(table, Table)

src/docx/blkcntnr.py

@@ -1,27 +1,48 @@
+# pyright: reportImportCycles=false
+
 """Block item container, used by body, cell, header, etc.
 
 Block level items are things like paragraph and table, although there are a few other
 specialized ones like structured document tags.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from typing_extensions import TypeAlias
+
 from docx.oxml.table import CT_Tbl
-from docx.shared import Parented
+from docx.shared import StoryChild
 from docx.text.paragraph import Paragraph
 
+if TYPE_CHECKING:
+ from docx import types as t
+ from docx.oxml.document import CT_Body
+ from docx.oxml.section import CT_HdrFtr
+ from docx.oxml.table import CT_Tc
+ from docx.shared import Length
+ from docx.styles.style import ParagraphStyle
+ from docx.table import Table
 
-class BlockItemContainer(Parented):
+BlockItemElement: TypeAlias = "CT_Body | CT_HdrFtr | CT_Tc"
+
+
+class BlockItemContainer(StoryChild):
  """Base class for proxy objects that can contain block items.
 
  These containers include _Body, _Cell, header, footer, footnote, endnote, comment,
  and text box objects. Provides the shared functionality to add a block item like a
  paragraph or table.
  """
 
- def __init__(self, element, parent):
+ def __init__(self, element: BlockItemElement, parent: t.ProvidesStoryPart):
  super(BlockItemContainer, self).__init__(parent)
  self._element = element
 
- def add_paragraph(self, text="", style=None):
+ def add_paragraph(
+ self, text: str = "", style: str | ParagraphStyle | None = None
+ ) -> Paragraph:
  """Return paragraph newly added to the end of the content in this container.
 
  The paragraph has `text` in a single run if present, and is given paragraph
@@ -37,7 +58,7 @@ def add_paragraph(self, text="", style=None):
  paragraph.style = style
  return paragraph
 
- def add_table(self, rows, cols, width):
+ def add_table(self, rows: int, cols: int, width: Length) -> Table:
  """Return table of `width` having `rows` rows and `cols` columns.
 
  The table is appended appended at the end of the content in this container.
@@ -47,7 +68,7 @@ def add_table(self, rows, cols, width):
  from docx.table import Table
 
  tbl = CT_Tbl.new_tbl(rows, cols, width)
- self._element._insert_tbl(tbl)
+ self._element._insert_tbl(tbl) # # pyright: ignore[reportPrivateUsage]
  return Table(tbl, self)
 
  @property
@@ -64,7 +85,7 @@ def tables(self):
 
  Read-only.
  """
- from .table import Table
+ from docx.table import Table
 
  return [Table(tbl, self) for tbl in self._element.tbl_lst]
 

src/docx/document.py

@@ -1,11 +1,27 @@
+# pyright: reportImportCycles=false
+# pyright: reportPrivateUsage=false
+
 """|Document| and closely related objects."""
 
+from __future__ import annotations
+
+from typing import IO, TYPE_CHECKING, List
+
 from docx.blkcntnr import BlockItemContainer
 from docx.enum.section import WD_SECTION
 from docx.enum.text import WD_BREAK
 from docx.section import Section, Sections
 from docx.shared import ElementProxy, Emu
-from docx.text.paragraph import Paragraph
+
+if TYPE_CHECKING:
+ from docx import types as t
+ from docx.oxml.document import CT_Body, CT_Document
+ from docx.parts.document import DocumentPart
+ from docx.settings import Settings
+ from docx.shared import Length
+ from docx.styles.style import ParagraphStyle, _TableStyle
+ from docx.table import Table
+ from docx.text.paragraph import Paragraph
 
 
 class Document(ElementProxy):
@@ -15,12 +31,13 @@ class Document(ElementProxy):
  a document.
  """
 
- def __init__(self, element, part):
+ def __init__(self, element: CT_Document, part: DocumentPart):
  super(Document, self).__init__(element)
+ self._element = element
  self._part = part
  self.__body = None
 
- def add_heading(self, text="", level=1):
+ def add_heading(self, text: str = "", level: int = 1):
  """Return a heading paragraph newly added to the end of the document.
 
  The heading paragraph will contain `text` and have its paragraph style
@@ -39,7 +56,9 @@ def add_page_break(self):
  paragraph.add_run().add_break(WD_BREAK.PAGE)
  return paragraph
 
- def add_paragraph(self, text: str = "", style=None) -> Paragraph:
+ def add_paragraph(
+ self, text: str = "", style: str | ParagraphStyle | None = None
+ ) -> Paragraph:
  """Return paragraph newly added to the end of the document.
 
  The paragraph is populated with `text` and having paragraph style `style`.
@@ -51,7 +70,12 @@ def add_paragraph(self, text: str = "", style=None) -> Paragraph:
  """
  return self._body.add_paragraph(text, style)
 
- def add_picture(self, image_path_or_stream, width=None, height=None):
+ def add_picture(
+ self,
+ image_path_or_stream: str | IO[bytes],
+ width: int | Length | None = None,
+ height: int | Length | None = None,
+ ):
  """Return new picture shape added in its own paragraph at end of the document.
 
  The picture contains the image at `image_path_or_stream`, scaled based on
@@ -65,7 +89,7 @@ def add_picture(self, image_path_or_stream, width=None, height=None):
  run = self.add_paragraph().add_run()
  return run.add_picture(image_path_or_stream, width, height)
 
- def add_section(self, start_type=WD_SECTION.NEW_PAGE):
+ def add_section(self, start_type: WD_SECTION = WD_SECTION.NEW_PAGE):
  """Return a |Section| object newly added at the end of the document.
 
  The optional `start_type` argument must be a member of the :ref:`WdSectionStart`
@@ -75,7 +99,7 @@ def add_section(self, start_type=WD_SECTION.NEW_PAGE):
  new_sectPr.start_type = start_type
  return Section(new_sectPr, self._part)
 
- def add_table(self, rows, cols, style=None):
+ def add_table(self, rows: int, cols: int, style: str | _TableStyle | None = None):
  """Add a table having row and column counts of `rows` and `cols` respectively.
 
  `style` may be a table style object or a table style name. If `style` is |None|,
@@ -92,7 +116,7 @@ def core_properties(self):
 
  @property
  def inline_shapes(self):
- """The |InlineShapes| collectoin for this document.
+ """The |InlineShapes| collection for this document.
 
  An inline shape is a graphical object, such as a picture, contained in a run of
  text and behaving like a character glyph, being flowed like other text in a
@@ -101,7 +125,7 @@ def inline_shapes(self):
  return self._part.inline_shapes
 
  @property
- def paragraphs(self):
+ def paragraphs(self) -> List[Paragraph]:
  """The |Paragraph| instances in the document, in document order.
 
  Note that paragraphs within revision marks such as ``<w:ins>`` or ``<w:del>`` do
@@ -110,11 +134,11 @@ def paragraphs(self):
  return self._body.paragraphs
 
  @property
- def part(self):
+ def part(self) -> DocumentPart:
  """The |DocumentPart| object of this document."""
  return self._part
 
- def save(self, path_or_stream):
+ def save(self, path_or_stream: str | IO[bytes]):
  """Save this document to `path_or_stream`.
 
  `path_or_stream` can be either a path to a filesystem location (a string) or a
@@ -123,12 +147,12 @@ def save(self, path_or_stream):
  self._part.save(path_or_stream)
 
  @property
- def sections(self):
+ def sections(self) -> Sections:
  """|Sections| object providing access to each section in this document."""
  return Sections(self._element, self._part)
 
  @property
- def settings(self):
+ def settings(self) -> Settings:
  """A |Settings| object providing access to the document-level settings."""
  return self._part.settings
 
@@ -138,7 +162,7 @@ def styles(self):
  return self._part.styles
 
  @property
- def tables(self):
+ def tables(self) -> List[Table]:
  """All |Table| instances in the document, in document order.
 
  Note that only tables appearing at the top level of the document appear in this
@@ -149,13 +173,13 @@ def tables(self):
  return self._body.tables
 
  @property
- def _block_width(self):
+ def _block_width(self) -> Length:
  """A |Length| object specifying the space between margins in last section."""
  section = self.sections[-1]
  return Emu(section.page_width - section.left_margin - section.right_margin)
 
  @property
- def _body(self):
+ def _body(self) -> _Body:
  """The |_Body| instance containing the content for this document."""
  if self.__body is None:
  self.__body = _Body(self._element.body, self)
@@ -168,7 +192,7 @@ class _Body(BlockItemContainer):
  It's primary role is a container for document content.
  """
 
- def __init__(self, body_elm, parent):
+ def __init__(self, body_elm: CT_Body, parent: t.ProvidesStoryPart):
  super(_Body, self).__init__(body_elm, parent)
  self._body = body_elm
 

src/docx/oxml/document.py

@@ -2,11 +2,15 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import TYPE_CHECKING, Callable, List
 
 from docx.oxml.section import CT_SectPr
 from docx.oxml.xmlchemy import BaseOxmlElement, ZeroOrMore, ZeroOrOne
 
+if TYPE_CHECKING:
+ from docx.oxml.table import CT_Tbl
+ from docx.oxml.text.paragraph import CT_P
+
 
 class CT_Document(BaseOxmlElement):
  """``<w:document>`` element, the root element of a document.xml file."""
@@ -29,14 +33,22 @@ def sectPr_lst(self) -> List[CT_SectPr]:
 
 
 class CT_Body(BaseOxmlElement):
- """``<w:body>``, the container element for the main document story in
- ``document.xml``."""
+ """`w:body`, the container element for the main document story in `document.xml`."""
+
+ add_p: Callable[[], CT_P]
+ get_or_add_sectPr: Callable[[], CT_SectPr]
+ p_lst: List[CT_P]
+ tbl_lst: List[CT_Tbl]
+
+ _insert_tbl: Callable[[CT_Tbl], CT_Tbl]
 
  p = ZeroOrMore("w:p", successors=("w:sectPr",))
  tbl = ZeroOrMore("w:tbl", successors=("w:sectPr",))
- sectPr = ZeroOrOne("w:sectPr", successors=())
+ sectPr: CT_SectPr | None = ZeroOrOne( # pyright: ignore[reportGeneralTypeIssues]
+ "w:sectPr", successors=()
+ )
 
- def add_section_break(self):
+ def add_section_break(self) -> CT_SectPr:
  """Return `w:sectPr` element for new section added at end of document.
 
  The last `w:sectPr` becomes the second-to-last, with the new `w:sectPr` being an
@@ -63,6 +75,5 @@ def clear_content(self):
 
  Leave the <w:sectPr> element if it is present.
  """
- content_elms = self[:-1] if self.sectPr is not None else self[:]
- for content_elm in content_elms:
+ for content_elm in self.xpath("./*[not(self::w:sectPr)]"):
  self.remove(content_elm)