Skip to content

Commit

Permalink
rfctr: improve typing local to BlockItemContainer
Browse files Browse the repository at this point in the history
  • Loading branch information
scanny committed Nov 3, 2023
1 parent cf17811 commit e315139
Show file tree
Hide file tree
Showing 11 changed files with 182 additions and 86 deletions.
13 changes: 7 additions & 6 deletions features/steps/block.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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()

Expand All @@ -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)

Expand All @@ -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)
35 changes: 28 additions & 7 deletions src/docx/blkcntnr.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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]

Expand Down
58 changes: 41 additions & 17 deletions src/docx/document.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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
Expand All @@ -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`.
Expand All @@ -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
Expand All @@ -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`
Expand All @@ -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|,
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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

Expand Down
25 changes: 18 additions & 7 deletions src/docx/oxml/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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
Expand All @@ -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)
Loading

0 comments on commit e315139

Please sign in to comment.