8000 rfctr: improve typing local to BlockItemContainer · python-openxml/python-docx@e315139 · GitHub
[go: up one dir, main page]

Skip to content

Commit e315139

Browse files
committed
rfctr: improve typing local to BlockItemContainer
1 parent cf17811 commit e315139

File tree

11 files changed

+182
-86
lines changed

11 files changed

+182
-86
lines changed

features/steps/block.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Step implementations for block content containers."""
22

33
from behave import given, then, when
4+
from behave.runner import Context
45

56
from docx import Document
67
from docx.table import Table
@@ -11,12 +12,12 @@
1112

1213

1314
@given("a document containing a table")
14-
def given_a_document_containing_a_table(context):
15+
def given_a_document_containing_a_table(context: Context):
1516
context.document = Document(test_docx("blk-containing-table"))
1617

1718

1819
@given("a paragraph")
19-
def given_a_paragraph(context):
20+
def given_a_paragraph(context: Context):
2021
context.document = Document()
2122
context.paragraph = context.document.add_paragraph()
2223

@@ -25,13 +26,13 @@ def given_a_paragraph(context):
2526

2627

2728
@when("I add a paragraph")
28-
def when_add_paragraph(context):
29+
def when_add_paragraph(context: Context):
2930
document = context.document
3031
context.p = document.add_paragraph()
3132

3233

3334
@when("I add a table")
34-
def when_add_table(context):
35+
def when_add_table(context: Context):
3536
rows, cols = 2, 2
3637
context.document.add_table(rows, cols)
3738

@@ -40,12 +41,12 @@ def when_add_table(context):
4041

4142

4243
@then("I can access the table")
43-
def then_can_access_table(context):
44+
def then_can_access_table(context: Context):
4445
table = context.document.tables[-1]
4546
assert isinstance(table, Table)
4647

4748

4849
@then("the new table appears in the document")
49-
def then_new_table_appears_in_document(context):
50+
def then_new_table_appears_in_document(context: Context):
5051
table = context.document.tables[-1]
5152
assert isinstance(table, Table)

src/docx/blkcntnr.py

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,48 @@
1+
# pyright: reportImportCycles=false
2+
13
"""Block item container, used by body, cell, header, etc.
24
35
Block level items are things like paragraph and table, although there are a few other
46
specialized ones like structured document tags.
57
"""
68

9+
from __future__ import annotations
10+
11+
from typing import TYPE_CHECKING
12+
13+
from typing_extensions import TypeAlias
14+
715
from docx.oxml.table import CT_Tbl
8-
from docx.shared import Parented
16+
from docx.shared import StoryChild
917
from docx.text.paragraph import Paragraph
1018

19+
if TYPE_CHECKING:
20+
from docx import types as t
21+
from docx.oxml.document import CT_Body
22+
from docx.oxml.section import CT_HdrFtr
23+
from docx.oxml.table import CT_Tc
24+
from docx.shared import Length
25+
from docx.styles.style import ParagraphStyle
26+
from docx.table import Table
1127

12-
class BlockItemContainer(Parented):
28+
BlockItemElement: TypeAlias = "CT_Body | CT_HdrFtr | CT_Tc"
29+
30+
31+
class BlockItemContainer(StoryChild):
1332
"""Base class for proxy objects that can contain block items.
1433
1534
These containers include _Body, _Cell, header, footer, footnote, endnote, comment,
1635
and text box objects. Provides the shared functionality to add a block item like a
1736
paragraph or table.
1837
"""
1938

20-
def __init__(self, element, parent):
39+
def __init__(self, element: BlockItemElement, parent: t.ProvidesStoryPart):
2140
super(BlockItemContainer, self).__init__(parent)
2241
self._element = element
2342

24-
def add_paragraph(self, text="", style=None):
43+
def add_paragraph(
44+
self, text: str = "", style: str | ParagraphStyle | None = None
45+
) -> Paragraph:
2546
"""Return paragraph newly added to the end of the content in this container.
2647
2748
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):
3758
paragraph.style = style
3859
return paragraph
3960

40-
def add_table(self, rows, cols, width):
61+
def add_table(self, rows: int, cols: int, width: Length) -> Table:
4162
"""Return table of `width` having `rows` rows and `cols` columns.
4263
4364
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):
4768
from docx.table import Table
4869

4970
tbl = CT_Tbl.new_tbl(rows, cols, width)
50-
self._element._insert_tbl(tbl)
71+
self._element._insert_tbl(tbl) # # pyright: ignore[reportPrivateUsage]
5172
return Table(tbl, self)
5273

5374
@property
@@ -64,7 +85,7 @@ def tables(self):
6485
6586
Read-only.
6687
"""
67-
from .table import Table
88+
from docx.table import Table
6889

6990
return [Table(tbl, self) for tbl in self._element.tbl_lst]
7091

src/docx/document.py

Lines changed: 41 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,27 @@
1+
# pyright: reportImportCycles=false
2+
# pyright: reportPrivateUsage=false
3+
14
"""|Document| and closely related objects."""
25

6+
from __future__ import annotations
7+
8+
from typing import IO, TYPE_CHECKING, List
9+
310
from docx.blkcntnr import BlockItemContainer
411
from docx.enum.section import WD_SECTION
512
from docx.enum.text import WD_BREAK
613
from docx.section import Section, Sections
714
from docx.shared import ElementProxy, Emu
8-
from docx.text.paragraph import Paragraph
15+
16+
if TYPE_CHECKING:
17+
from docx import types as t
18+
from docx.oxml.document import CT_Body, CT_Document
19+
from docx.parts.document import DocumentPart
20+
from docx.settings import Settings
21+
from docx.shared import Length
22+
from docx.styles.style import ParagraphStyle, _TableStyle
23+
from docx.table import Table
24+
from docx.text.paragraph import Paragraph
925

1026

1127
class Document(ElementProxy):
@@ -15,12 +31,13 @@ class Document(ElementProxy):
1531
a document.
1632
"""
1733

18-
def __init__(self, element, part):
34+
def __init__(self, element: CT_Document, part: DocumentPart):
1935
super(Document, self).__init__(element)
36+
self._element = element
2037
self._part = part
2138
self.__body = None
2239

23-
def add_heading(self, text="", level=1):
40+
def add_heading(self, text: str = "", level: int = 1):
2441
"""Return a heading paragraph newly added to the end of the document.
2542
2643
The heading paragraph will contain `text` and have its paragraph style
@@ -39,7 +56,9 @@ def add_page_break(self):
3956
paragraph.add_run().add_break(WD_BREAK.PAGE)
4057
return paragraph
4158

42-
def add_paragraph(self, text: str = "", style=None) -> Paragraph:
59+
def add_paragraph(
60+
self, text: str = "", style: str | ParagraphStyle | None = None
61+
) -> Paragraph:
4362
"""Return paragraph newly added to the end of the document.
4463
4564
The paragraph is populated with `text` and having paragraph style `style`.
@@ -51,7 +70,12 @@ def add_paragraph(self, text: str = "", style=None) -> Paragraph:
5170
"""
5271
return self._body.add_paragraph(text, style)
5372

54-
def add_picture(self, image_path_or_stream, width=None, height=None):
73+
def add_picture(
74+
self,
75+
image_path_or_stream: str | IO[bytes],
76+
width: int | Length | None = None,
77+
height: int | Length | None = None,
78+
):
5579
"""Return new picture shape added in its own paragraph at end of the document.
5680
5781
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):
6589
run = self.add_paragraph().add_run()
6690
return run.add_picture(image_path_or_stream, width, height)
6791

68-
def add_section(self, start_type=WD_SECTION.NEW_PAGE):
92+
def add_section(self, start_type: WD_SECTION = WD_SECTION.NEW_PAGE):
6993
"""Return a |Section| object newly added at the end of the document.
7094
7195
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):
7599
new_sectPr.start_type = start_type
76100
return Section(new_sectPr, self._part)
77101

78-
def add_table(self, rows, cols, style=None):
102+
def add_table(self, rows: int, cols: int, style: str | _TableStyle | None = None):
79103
"""Add a table having row and column counts of `rows` and `cols` respectively.
80104
81105
`style` may be a table style object or a table style name. If `style` is |None|,
@@ -92,7 +116,7 @@ def core_properties(self):
92116

93117
@property
94118
def inline_shapes(self):
95-
"""The |InlineShapes| collectoin for this document.
119+
"""The |InlineShapes| collection for this document.
96120
97121
An inline shape is a graphical object, such as a picture, contained in a run of
98122
text and behaving like a character glyph, being flowed like other text in a
@@ -101,7 +125,7 @@ def inline_shapes(self):
101125
return self._part.inline_shapes
102126

103127
@property
104-
def paragraphs(self):
128+
def paragraphs(self) -> List[Paragraph]:
105129
"""The |Paragraph| instances in the document, in document order.
106130
107131
Note that paragraphs within revision marks such as ``<w:ins>`` or ``<w:del>`` do
@@ -110,11 +134,11 @@ def paragraphs(self):
110134
return self._body.paragraphs
111135

112136
@property
113-
def part(self):
137+
def part(self) -> DocumentPart:
114138
"""The |DocumentPart| object of this document."""
115139
return self._part
116140

117-
def save(self, path_or_stream):
141+
def save(self, path_or_stream: str | IO[bytes]):
118142
"""Save this document to `path_or_stream`.
119143
120144
`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):
123147
self._part.save(path_or_stream)
124148

125149
@property
126-
def sections(self):
150+
def sections(self) -> Sections:
127151
"""|Sections| object providing access to each section in this document."""
128152
return Sections(self._element, self._part)
129153

130154
@property
131-
def settings(self):
155+
def settings(self) -> Settings:
132156
"""A |Settings| object providing access to the document-level settings."""
133157
return self._part.settings
134158

@@ -138,7 +162,7 @@ def styles(self):
138162
return self._part.styles
139163

140164
@property
141-
def tables(self):
165+
def tables(self) -> List[Table]:
142166
"""All |Table| instances in the document, in document order.
143167
144168
Note that only tables appearing at the top level of the document appear in this
@@ -149,13 +173,13 @@ def tables(self):
149173
return self._body.tables
150174

151175
@property
152-
def _block_width(self):
176+
def _block_width(self) -> Length:
153177
"""A |Length| object specifying the space between margins in last section."""
154178
section = self.sections[-1]
155179
return Emu(section.page_width - section.left_margin - section.right_margin)
156180

157181
@property
158-
def _body(self):
182+
def _body(self) -> _Body:
159183
"""The |_Body| instance containing the content for this document."""
160184
if self.__body is None:
161185
self.__body = _Body(self._element.body, self)
@@ -168,7 +192,7 @@ class _Body(BlockItemContainer):
168192
It's primary role is a container for document content.
169193
"""
170194

171-
def __init__(self, body_elm, parent):
195+
def __init__(self, body_elm: CT_Body, parent: t.ProvidesStoryPart):
172196
super(_Body, self).__init__(body_elm, parent)
173197
self._body = body_elm
174198

src/docx/oxml/document.py

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,15 @@
22

33
from __future__ import annotations
44

5-
from typing import List
5+
from typing import TYPE_CHECKING, Callable, List
66

77
from docx.oxml.section import CT_SectPr
88
from docx.oxml.xmlchemy import BaseOxmlElement, ZeroOrMore, ZeroOrOne
99

10+
if TYPE_CHECKING:
11+
from docx.oxml.table import CT_Tbl
12+
from docx.oxml.text.paragraph import CT_P
13+
1014

1115
class CT_Document(BaseOxmlElement):
1216
"""``<w:document>`` element, the root element of a document.xml file."""
@@ -29,14 +33,22 @@ def sectPr_lst(self) -> List[CT_SectPr]:
2933

3034

3135
class CT_Body(BaseOxmlElement):
32-
"""``<w:body>``, the container element for the main document story in
33-
``document.xml``."""
36+
"""`w:body`, the container element for the main document story in `document.xml`."""
37+
38+
add_p: Callable[[], CT_P]
39+
get_or_add_sectPr: Callable[[], CT_SectPr]
40+
p_lst: List[CT_P]
41+
tbl_lst: List[CT_Tbl]
42+
43+
_insert_tbl: Callable[[CT_Tbl], CT_Tbl]
3444

3545
p = ZeroOrMore("w:p", successors=("w:sectPr",))
3646
tbl = ZeroOrMore("w:tbl", successors=("w:sectPr",))
37-
sectPr = ZeroOrOne("w:sectPr", successors=())
47+
sectPr: CT_SectPr | None = ZeroOrOne( # pyright: ignore[reportGeneralTypeIssues]
48+
"w:sectPr", successors=()
49+
)
3850

39-
def add_section_break(self):
51+
def add_section_break(self) -> CT_SectPr:
4052
"""Return `w:sectPr` element for new section added at end of document.
4153
4254
The last `w:sectPr` becomes the second-to-last, with the new `w:sectPr` being an
@@ -63,6 +75,5 @@ def clear_content(self):
6375
6476
Leave the <w:sectPr> element if it is present.
6577
"""
66-
content_elms = self[:-1] if self.sectPr is not None else self[:]
67-
for content_elm in content_elms:
78+
for content_elm in self.xpath("./*[not(self::w:sectPr)]"):
6879
self.remove(content_elm)

0 commit comments

Comments
 (0)
0