36646375051

Committed 29 Sep 2026 11:39PM UTC coverage: 90.867% (-0.06%) from 90.925%

Build # 36646375051

Build Type

push

github

Committed by

web-flow

Commit Message

Added type annotations and fixed mypy errors in python-ta package (#1397)

Coverage Stats

256 of 269 new or added lines in 33 files covered. (95.17%)

2 existing lines in 2 files now uncovered.

3761 of 4139 relevant lines covered (90.87%)

17.67 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.0

/packages/python-ta/src/python_ta/transforms/setendings.py

"""
Top-level functions to mutate the astroid nodes with `end_col_offset` and
`end_lineno` properties.

Where possible, the `end_col_offset` property is set by that of the node's last child.

    fromlineno
        - existing attribute
        - one-indexed
    end_lineno
        - new attribute
        - one-indexed
    col_offset
        - existing attribute
        - zero-indexed
        - located left of the first character
    end_col_offset
        - new attribute
        - zero-indexed
        - located right of the last character (essentially the string length)

In astroid/astroid/transforms.py, functions are registered to types in the
`transforms` dictionary in the TransformVisitor class. The traversal at
line 83 eventually leads to the transform called on each node at line 36,
within the _transform method.

Astroid Source:
https://github.com/PyCQA/astroid/blob/master/astroid/transforms.py
"""

from collections.abc import Callable
from typing import Any

from astroid import nodes
from astroid.transforms import TransformVisitor
from pylint.lint import PyLinter

NodePredicate = Callable[[str, int, nodes.NodeNG], bool]
NodeTransform = Callable[[nodes.NodeNG], nodes.NodeNG]

CONSUMABLES = " \n\t\\"

# These nodes have a child, and their end_lineno and end_col_offset
# attributes are set equal to those of their last child.
NODES_WITH_CHILDREN = [
    nodes.Call,
    nodes.Comprehension,
    nodes.Module,
]


# Predicate functions, for setting locations based on source code.
# Predicates can only return a single truthy value, because of how its used in
# `astroid/transforms.py`
# ====================================================
def _token_search(token: str) -> NodePredicate:
    """
    @type token: string
    @rtype: function
    """

    def _is_token(s: str, index: int, node: nodes.NodeNG) -> bool:
        """Fix to include certain tokens such as a paren, bracket, or brace.
        @type s: string
        @type index: int
        @type node: Astroid node
        @rtype: bool
        """
        return s[index] == token

    return _is_token


def _keyword_search(keyword: str) -> NodePredicate:
    """
    @type keyword: string
    @rtype: function
    """

    def _is_keyword(s: str, index: int, node: nodes.NodeNG) -> bool:
        """Search for a keyword. Right-to-left.
        @type s: string
        @type index: int
        @type node: Astroid node
        @rtype: bool
        """
        return s[index : index + len(keyword)] == keyword

    return _is_keyword


def _is_arg_name(s: str, index: int, node: nodes.NodeNG) -> bool:
    """Search for the name of the argument. Right-to-left."""
    if not node.arg:
        return False
    return s[index : index + len(node.arg)] == node.arg


# Nodes the require the source code for proper location setting
# Elements here are in the form
# (node class, predicate for start | None, predicate for end | None)
NODES_REQUIRING_SOURCE = [
    (nodes.Call, None, _token_search(")")),
    (nodes.DelAttr, _keyword_search("del"), None),
    (nodes.DelName, _keyword_search("del"), None),
    (nodes.GeneratorExp, _token_search("("), _token_search(")")),
    (nodes.List, _token_search("["), _token_search("]")),
    (nodes.Tuple, None, _token_search(",")),
]


def init_register_ending_setters(source_code: list[str]) -> TransformVisitor:
    """Instantiate a visitor to transform the nodes.
    Register the transform functions on an instance of TransformVisitor.

    @type source_code: list of strings
    @rtype: TransformVisitor
    """
    ending_transformer = TransformVisitor()

    # Check consistency of astroid-provided fromlineno and col_offset attributes.
    for node_class in nodes.ALL_NODE_CLASSES:
        ending_transformer.register_transform(
            node_class,
            fix_start_attributes,
            lambda node: (
                getattr(node, "fromlineno", None) is None
                or getattr(node, "col_offset", None) is None
            ),
        )

    # Ad hoc transformations
    ending_transformer.register_transform(nodes.BinOp, _set_start_from_first_child)
    ending_transformer.register_transform(nodes.ClassDef, _set_start_from_first_decorator)
    ending_transformer.register_transform(nodes.FunctionDef, _set_start_from_first_decorator)
    ending_transformer.register_transform(nodes.Tuple, _set_start_from_first_child)
    ending_transformer.register_transform(nodes.Arguments, fix_arguments(source_code))

    for node_class in NODES_WITH_CHILDREN:
        ending_transformer.register_transform(node_class, set_from_last_child)

    # Nodes where the source code must also be provided.
    # source_code and the predicate functions get stored in the TransformVisitor
    for node_class, start_pred, end_pred in NODES_REQUIRING_SOURCE:
        if start_pred is not None:
            ending_transformer.register_transform(
                node_class, start_setter_from_source(source_code, start_pred)
            )
        if end_pred is not None:
            # This is for searching for a trailing comma after a tuple's final element
            if node_class is nodes.Tuple:
                ending_transformer.register_transform(
                    node_class, end_setter_from_source(source_code, end_pred, True)
                )
            else:
                ending_transformer.register_transform(
                    node_class, end_setter_from_source(source_code, end_pred)
                )

    # Nodes where extra parentheses are included
    ending_transformer.register_transform(nodes.BinOp, add_parens(source_code))
    ending_transformer.register_transform(nodes.Const, add_parens(source_code))
    ending_transformer.register_transform(nodes.Tuple, add_parens(source_code))

    return ending_transformer


# Transform functions.
# These functions are called on individual nodes to either fix the
# `fromlineno` and `col_offset` properties of the nodes,
# or to set the `end_lineno` and `end_col_offset` attributes for a node.
# ====================================================
def fix_arguments(source_code: list[str]) -> NodeTransform:
    """For an Arguments node"""

    def _find(node: nodes.Arguments) -> nodes.Arguments:
        children = list(node.get_children())
        if children:
            fix_start_attributes(node)

        line_i = node.parent.fromlineno
        char_i = node.parent.col_offset
        for child in children:
            if line_i is None:
                line_i = child.end_lineno
                char_i = child.end_col_offset
            elif (
                line_i < child.end_lineno
                or line_i == child.end_lineno
                and char_i < child.end_col_offset
            ):
                line_i = child.end_lineno
                char_i = child.end_col_offset

        line_i -= 1  # Switch to 0-indexing

        # left bracket if parent is FunctionDef, colon if Lambda
        if isinstance(node.parent, nodes.FunctionDef):
            end_char = ")"
        else:
            end_char = ":"

        while char_i < len(source_code[line_i]) and source_code[line_i][char_i] != end_char:
            if char_i == len(source_code[line_i]) - 1 or source_code[line_i][char_i] == "#":
                char_i = 0
                line_i += 1
            else:
                char_i += 1

        node.end_lineno, node.end_col_offset = line_i + 1, char_i

        # no children
        if children == []:
            node.fromlineno, node.col_offset = line_i + 1, char_i

        return node

    return _find


def fix_start_attributes(node: nodes.NodeNG) -> nodes.NodeNG:
    """Some nodes don't always have the `col_offset` property set by Astroid:
    Comprehension, Keyword, Module, Slice.
    """
    try:
        first_child = next(node.get_children())
        if getattr(node, "fromlineno", None) is None:
            node.fromlineno = first_child.fromlineno
        if getattr(node, "col_offset", None) is None:
            node.col_offset = first_child.col_offset

    except StopIteration:
        # No children. Go to the enclosing statement and use that.
        # This assumes that statement nodes will always have these attributes set.
        statement = node.statement()
        if statement is not node:
            if getattr(node, "fromlineno", None) is None:
                node.fromlineno = statement.fromlineno
            if getattr(node, "col_offset", None) is None:
                node.col_offset = statement.col_offset
        else:
            # Enclosing statement is same as node, also does not have attributes set
            if getattr(node, "fromlineno", None) is None:
                node.fromlineno = 0
            if getattr(node, "col_offset", None) is None:
                node.col_offset = 0
    return node


def _set_start_from_first_child(node: nodes.NodeNG) -> nodes.NodeNG:
    """Set the start attributes of this node from its first child."""
    try:
        first_child = next(node.get_children())
    except StopIteration:
        pass
    else:
        node.fromlineno = first_child.fromlineno
        node.col_offset = first_child.col_offset
    return node


def _set_start_from_first_decorator(node: nodes.NodeNG) -> nodes.NodeNG:
    """Set the start attributes of this node from its first child, if that child is a decorator."""
    if getattr(node, "decorators"):
        first_child = node.decorators
        node.fromlineno = first_child.fromlineno
        node.col_offset = first_child.col_offset
    return node


def set_from_last_child(node: nodes.NodeNG) -> nodes.NodeNG:
    """Populate ending locations for astroid node based on its last child.

    Preconditions:
      - `node` must have a `last_child` (node).
      - `node` has col_offset property set.
    """
    last_child = _get_last_child(node)
    if not last_child:
        return node

    if last_child.end_lineno is not None:
        node.end_lineno = last_child.end_lineno
    if last_child.end_col_offset is not None:
        node.end_col_offset = last_child.end_col_offset
    return node


def _get_last_child(node: nodes.NodeNG) -> nodes.NodeNG | None:
    """Returns the last child node, or None.
    Some nodes' last_child() attribute not set, e.g. nodes.Arguments.
    """
    if node.last_child():
        return node.last_child()
    else:
        # Get the first child from the `get_children` generator.
        skip_to_last_child = None  # save reference to last child.
        for skip_to_last_child in node.get_children():
            pass  # skip to last
        return skip_to_last_child  # postcondition: node, or None.


def end_setter_from_source(
    source_code: list[str], pred: NodePredicate, only_consumables: bool = False
) -> NodeTransform:
    """Returns a *function* that sets ending locations for a node from source.

    The basic technique is to do the following:
      1. Find the ending locations for the node based on its last child.
      2. Starting at that point, iterate through characters in the source code
         up to and including the first index that satisfies pred.

    pred is a function that takes a string and index and returns a bool,
    e.g. _is_close_paren

    If only_consumables is True, the search halts when it reaches a non-consumable
    character that fails pred *on the first line*.
    TODO: really the behaviour should be the same for all lines searched for.
    """

    def set_endings_from_source(node: nodes.NodeNG) -> nodes.NodeNG:
        # Tuple nodes have an end_col_offset that includes the end paren,
        # but their col_offset does not include the start paren.
        # To address this, we override the Tuple node's end_col_offset.
        if isinstance(node, nodes.Tuple):
            set_from_last_child(node)

        # Initialize counters. Note: we need to offset lineno,
        # since it's 1-indexed.
        end_col_offset, lineno = node.end_col_offset, node.end_lineno - 1

        # First, search the remaining part of the current end line.
        for j in range(end_col_offset, len(source_code[lineno])):
            if source_code[lineno][j] == "#":
                break  # skip over comment lines
            if pred(source_code[lineno], j, node):
                node.end_col_offset = j + 1
                return node
            elif only_consumables and source_code[lineno][j] not in CONSUMABLES:
                return node

        # If that doesn't work, search remaining lines
        for i in range(lineno + 1, len(source_code)):
            # Search each character
            for j in range(len(source_code[i])):
                if source_code[i][j] == "#":
                    break  # skip over comment lines
                if pred(source_code[i], j, node):
                    node.end_col_offset, node.end_lineno = j + 1, i + 1
                    return node
                # only consume inert characters.
                elif source_code[i][j] not in CONSUMABLES:
                    return node
        return node

    return set_endings_from_source


def start_setter_from_source(source_code: list[str], pred: NodePredicate) -> NodeTransform:
    """Returns a *function* that sets start locations for a node from source.
    Recall `source_code`, `pred` are within the lexical scope of the returned function.

    The basic technique is to do the following:
      1. Find the start locations for the node (already set).
      2. Starting at that point, iterate through characters in the source code
         in reverse until reaching the first index that satisfies pred.

    pred is a function that takes a string and index and returns a bool,
    e.g. _is_open_paren
    """

    def set_start_from_source(node: nodes.NodeNG) -> nodes.NodeNG:
        # Initialize counters. Note: fromlineno is 1-indexed.
        col_offset, lineno = node.col_offset, node.fromlineno - 1

        # First, search the remaining part of the current start line
        for j in range(min(len(source_code[lineno]) - 1, col_offset), -1, -1):
            if pred(source_code[lineno], j, node):
                node.col_offset = j
                return node

        # If that doesn't work, search remaining lines
        for i in range(lineno - 1, -1, -1):
            # Search each character, right-to-left
            for j in range(len(source_code[i]) - 1, -1, -1):
                if pred(source_code[i], j, node):
                    node.end_col_offset, node.end_lineno = j, i + 1
                    return node
                # only consume inert characters.
                elif source_code[i][j] not in CONSUMABLES:
                    return node
        return node

    return set_start_from_source


def add_parens(source_code: list[str]) -> NodeTransform:
    def h(node: nodes.NodeNG) -> nodes.NodeNG:
        _add_parens(source_code)(node)

    return h


def _add_parens(source_code: list[str]) -> NodeTransform:
    def h(node: nodes.NodeNG) -> nodes.NodeNG:
        # Initialize counters. Note: fromlineno is 1-indexed.
        prev = node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset
        while True:
            col_offset, lineno = node.col_offset, node.fromlineno - 1
            end_col_offset, end_lineno = node.end_col_offset, node.end_lineno - 1

            # First, search the remaining part of the current start line
            prev_char, new_lineno, new_coloffset = None, None, None
            for j in range(col_offset - 1, -1, -1):
                if source_code[lineno][j] in CONSUMABLES or source_code[lineno][j] == ",":
                    continue
                else:
                    prev_char, new_lineno, new_coloffset = source_code[lineno][j], lineno, j
                    break

            if prev_char is None:
                # Search remaining lines
                for i in range(lineno - 1, -1, -1):
                    # Search each character, right-to-left
                    for j in range(len(source_code[i]) - 1, -1, -1):
                        if source_code[i][j] in CONSUMABLES or source_code[i][j] == ",":
                            continue
                        else:
                            prev_char, new_lineno, new_coloffset = source_code[i][j], i, j

                            break
                    if prev_char is not None:
                        break

            if prev_char is None or new_lineno is None or new_coloffset is None:
                break

            if prev_char != "(":
                # No enclosing parentheses
                break

            # Now search for matching ')'
            next_char, new_end_lineno, new_end_coloffset = None, None, None
            for j in range(end_col_offset, len(source_code[end_lineno])):
                if source_code[end_lineno][j] == "#":
                    break  # skip over comment lines
                elif source_code[end_lineno][j] in CONSUMABLES:
                    continue
                else:
                    next_char, new_end_lineno, new_end_coloffset = (
                        source_code[end_lineno][j],
                        end_lineno,
                        j,
                    )
                    break

            if next_char is None:
                # Search remaining lines
                for i in range(end_lineno + 1, len(source_code)):
                    # Search each character
                    for j in range(len(source_code[i])):
                        if source_code[i][j] == "#":
                            break  # skip over comment lines
                        elif source_code[i][j] in CONSUMABLES:
                            continue
                        else:
                            next_char, new_end_lineno, new_end_coloffset = source_code[i][j], i, j
                            break
                    if next_char is not None:
                        break

            if next_char is None or new_end_lineno is None or new_end_coloffset is None:
                break

            if next_char != ")":
                break

            # At this point, an enclosing pair of parentheses has been found
            prev = node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset
            node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset = (
                new_lineno + 1,
                new_coloffset,
                new_end_lineno + 1,
                new_end_coloffset + 1,
            )

        # Go back by 1 set of parentheses if inside a function call.
        if isinstance(node.parent, nodes.Call) and len(node.parent.args) == 1:
            node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset = prev

        return node

    return h


# Make this module a pylint plugin
def register(linter: PyLinter) -> None:
    """Patch linter to apply message transform with source code."""
    old_get_ast = linter.get_ast

    def new_get_ast(filepath: str, modname: str, data: str) -> Any:
        ast = old_get_ast(filepath, modname, data)
        if ast is not None:
            with open(filepath, encoding="utf-8") as f:
                source_code = f.readlines()
            ending_transformer = init_register_ending_setters(source_code)
            ending_transformer.visit(ast)
        return ast

    linter.get_ast = new_get_ast  # type: ignore[assignment, method-assign]

1	"""
2	Top-level functions to mutate the astroid nodes with `end_col_offset` and
3	`end_lineno` properties.
4
5	Where possible, the `end_col_offset` property is set by that of the node's last child.
6
7	fromlineno
8	- existing attribute
9	- one-indexed
10	end_lineno
11	- new attribute
12	- one-indexed
13	col_offset
14	- existing attribute
15	- zero-indexed
16	- located left of the first character
17	end_col_offset
18	- new attribute
19	- zero-indexed
20	- located right of the last character (essentially the string length)
21
22	In astroid/astroid/transforms.py, functions are registered to types in the
23	`transforms` dictionary in the TransformVisitor class. The traversal at
24	line 83 eventually leads to the transform called on each node at line 36,
25	within the _transform method.
26
27	Astroid Source:
28	https://github.com/PyCQA/astroid/blob/master/astroid/transforms.py
29	"""
30
31	from collections.abc import Callable	20 ✔
32	from typing import Any	20 ✔
33
34	from astroid import nodes	20 ✔
35	from astroid.transforms import TransformVisitor	20 ✔
36	from pylint.lint import PyLinter	20 ✔
37
38	NodePredicate = Callable[[str, int, nodes.NodeNG], bool]	20 ✔
39	NodeTransform = Callable[[nodes.NodeNG], nodes.NodeNG]	20 ✔
40
41	CONSUMABLES = " \n\t\\"	20 ✔
42
43	# These nodes have a child, and their end_lineno and end_col_offset
44	# attributes are set equal to those of their last child.
45	NODES_WITH_CHILDREN = [	20 ✔
46	nodes.Call,
47	nodes.Comprehension,
48	nodes.Module,
49	]
50
51
52	# Predicate functions, for setting locations based on source code.
53	# Predicates can only return a single truthy value, because of how its used in
54	# `astroid/transforms.py`
55	# ====================================================
56	def _token_search(token: str) -> NodePredicate:	20 ✔
57	"""
58	@type token: string
59	@rtype: function
60	"""
61
62	def _is_token(s: str, index: int, node: nodes.NodeNG) -> bool:	20 ✔
63	"""Fix to include certain tokens such as a paren, bracket, or brace.
64	@type s: string
65	@type index: int
66	@type node: Astroid node
67	@rtype: bool
68	"""
69	return s[index] == token	20 ✔
70
71	return _is_token	20 ✔
72
73
74	def _keyword_search(keyword: str) -> NodePredicate:	20 ✔
75	"""
76	@type keyword: string
77	@rtype: function
78	"""
79
80	def _is_keyword(s: str, index: int, node: nodes.NodeNG) -> bool:	20 ✔
81	"""Search for a keyword. Right-to-left.
82	@type s: string
83	@type index: int
84	@type node: Astroid node
85	@rtype: bool
86	"""
87	return s[index : index + len(keyword)] == keyword	20 ✔
88
89	return _is_keyword	20 ✔
90
91
92	def _is_arg_name(s: str, index: int, node: nodes.NodeNG) -> bool:	20 ✔
93	"""Search for the name of the argument. Right-to-left."""
94	if not node.arg:	×
95	return False	×
96	return s[index : index + len(node.arg)] == node.arg	×
97
98
99	# Nodes the require the source code for proper location setting
100	# Elements here are in the form
101	# (node class, predicate for start \| None, predicate for end \| None)
102	NODES_REQUIRING_SOURCE = [	20 ✔
103	(nodes.Call, None, _token_search(")")),
104	(nodes.DelAttr, _keyword_search("del"), None),
105	(nodes.DelName, _keyword_search("del"), None),
106	(nodes.GeneratorExp, _token_search("("), _token_search(")")),
107	(nodes.List, _token_search("["), _token_search("]")),
108	(nodes.Tuple, None, _token_search(",")),
109	]
110
111
112	def init_register_ending_setters(source_code: list[str]) -> TransformVisitor:	20 ✔
113	"""Instantiate a visitor to transform the nodes.
114	Register the transform functions on an instance of TransformVisitor.
115
116	@type source_code: list of strings
117	@rtype: TransformVisitor
118	"""
119	ending_transformer = TransformVisitor()	20 ✔
120
121	# Check consistency of astroid-provided fromlineno and col_offset attributes.
122	for node_class in nodes.ALL_NODE_CLASSES:	20 ✔
123	ending_transformer.register_transform(	20 ✔
124	node_class,
125	fix_start_attributes,
126	lambda node: (
127	getattr(node, "fromlineno", None) is None
128	or getattr(node, "col_offset", None) is None
129	),
130	)
131
132	# Ad hoc transformations
133	ending_transformer.register_transform(nodes.BinOp, _set_start_from_first_child)	20 ✔
134	ending_transformer.register_transform(nodes.ClassDef, _set_start_from_first_decorator)	20 ✔
135	ending_transformer.register_transform(nodes.FunctionDef, _set_start_from_first_decorator)	20 ✔
136	ending_transformer.register_transform(nodes.Tuple, _set_start_from_first_child)	20 ✔
137	ending_transformer.register_transform(nodes.Arguments, fix_arguments(source_code))	20 ✔
138
139	for node_class in NODES_WITH_CHILDREN:	20 ✔
140	ending_transformer.register_transform(node_class, set_from_last_child)	20 ✔
141
142	# Nodes where the source code must also be provided.
143	# source_code and the predicate functions get stored in the TransformVisitor
144	for node_class, start_pred, end_pred in NODES_REQUIRING_SOURCE:	20 ✔
145	if start_pred is not None:	20 ✔
146	ending_transformer.register_transform(	20 ✔
147	node_class, start_setter_from_source(source_code, start_pred)
148	)
149	if end_pred is not None:	20 ✔
150	# This is for searching for a trailing comma after a tuple's final element
151	if node_class is nodes.Tuple:	20 ✔
152	ending_transformer.register_transform(	20 ✔
153	node_class, end_setter_from_source(source_code, end_pred, True)
154	)
155	else:
156	ending_transformer.register_transform(	20 ✔
157	node_class, end_setter_from_source(source_code, end_pred)
158	)
159
160	# Nodes where extra parentheses are included
161	ending_transformer.register_transform(nodes.BinOp, add_parens(source_code))	20 ✔
162	ending_transformer.register_transform(nodes.Const, add_parens(source_code))	20 ✔
163	ending_transformer.register_transform(nodes.Tuple, add_parens(source_code))	20 ✔
164
165	return ending_transformer	20 ✔
166
167
168	# Transform functions.
169	# These functions are called on individual nodes to either fix the
170	# `fromlineno` and `col_offset` properties of the nodes,
171	# or to set the `end_lineno` and `end_col_offset` attributes for a node.
172	# ====================================================
173	def fix_arguments(source_code: list[str]) -> NodeTransform:	20 ✔
174	"""For an Arguments node"""
175
176	def _find(node: nodes.Arguments) -> nodes.Arguments:	20 ✔
177	children = list(node.get_children())	20 ✔
178	if children:	20 ✔
179	fix_start_attributes(node)	20 ✔
180
181	line_i = node.parent.fromlineno	20 ✔
182	char_i = node.parent.col_offset	20 ✔
183	for child in children:	20 ✔
184	if line_i is None:	20 ✔
185	line_i = child.end_lineno	×
186	char_i = child.end_col_offset	×
187	elif (	20 ✔
188	line_i < child.end_lineno
189	or line_i == child.end_lineno
190	and char_i < child.end_col_offset
191	):
192	line_i = child.end_lineno	20 ✔
193	char_i = child.end_col_offset	20 ✔
194
195	line_i -= 1 # Switch to 0-indexing	20 ✔
196
197	# left bracket if parent is FunctionDef, colon if Lambda
198	if isinstance(node.parent, nodes.FunctionDef):	20 ✔
199	end_char = ")"	20 ✔
200	else:
201	end_char = ":"	20 ✔
202
203	while char_i < len(source_code[line_i]) and source_code[line_i][char_i] != end_char:	20 ✔
204	if char_i == len(source_code[line_i]) - 1 or source_code[line_i][char_i] == "#":	20 ✔
205	char_i = 0	20 ✔
206	line_i += 1	20 ✔
207	else:
208	char_i += 1	20 ✔
209
210	node.end_lineno, node.end_col_offset = line_i + 1, char_i	20 ✔
211
212	# no children
213	if children == []:	20 ✔
214	node.fromlineno, node.col_offset = line_i + 1, char_i	20 ✔
215
216	return node	20 ✔
217
218	return _find	20 ✔
219
220
221	def fix_start_attributes(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
222	"""Some nodes don't always have the `col_offset` property set by Astroid:
223	Comprehension, Keyword, Module, Slice.
224	"""
225	try:	20 ✔
226	first_child = next(node.get_children())	20 ✔
227	if getattr(node, "fromlineno", None) is None:	20 ✔
228	node.fromlineno = first_child.fromlineno	×
229	if getattr(node, "col_offset", None) is None:	20 ✔
230	node.col_offset = first_child.col_offset	20 ✔
231
232	except StopIteration:	20 ✔
233	# No children. Go to the enclosing statement and use that.
234	# This assumes that statement nodes will always have these attributes set.
235	statement = node.statement()	20 ✔
236	if statement is not node:	20 ✔
237	if getattr(node, "fromlineno", None) is None:	20 ✔
238	node.fromlineno = statement.fromlineno	×
239	if getattr(node, "col_offset", None) is None:	20 ✔
240	node.col_offset = statement.col_offset	20 ✔
241	else:
242	# Enclosing statement is same as node, also does not have attributes set
243	if getattr(node, "fromlineno", None) is None:	×
244	node.fromlineno = 0	×
245	if getattr(node, "col_offset", None) is None:	×
246	node.col_offset = 0	×
247	return node	20 ✔
248
249
250	def _set_start_from_first_child(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
251	"""Set the start attributes of this node from its first child."""
252	try:	20 ✔
253	first_child = next(node.get_children())	20 ✔
254	except StopIteration:	20 ✔
255	pass	20 ✔
256	else:
257	node.fromlineno = first_child.fromlineno	20 ✔
258	node.col_offset = first_child.col_offset	20 ✔
259	return node	20 ✔
260
261
262	def _set_start_from_first_decorator(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
263	"""Set the start attributes of this node from its first child, if that child is a decorator."""
264	if getattr(node, "decorators"):	20 ✔
265	first_child = node.decorators	20 ✔
266	node.fromlineno = first_child.fromlineno	20 ✔
267	node.col_offset = first_child.col_offset	20 ✔
268	return node	20 ✔
269
270
271	def set_from_last_child(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
272	"""Populate ending locations for astroid node based on its last child.
273
274	Preconditions:
275	- `node` must have a `last_child` (node).
276	- `node` has col_offset property set.
277	"""
278	last_child = _get_last_child(node)	20 ✔
279	if not last_child:	20 ✔
280	return node	20 ✔
281
282	if last_child.end_lineno is not None:	20 ✔
283	node.end_lineno = last_child.end_lineno	20 ✔
284	if last_child.end_col_offset is not None:	20 ✔
285	node.end_col_offset = last_child.end_col_offset	20 ✔
286	return node	20 ✔
287
288
289	def _get_last_child(node: nodes.NodeNG) -> nodes.NodeNG \| None:	20 ✔
290	"""Returns the last child node, or None.
291	Some nodes' last_child() attribute not set, e.g. nodes.Arguments.
292	"""
293	if node.last_child():	20 ✔
294	return node.last_child()	20 ✔
295	else:
296	# Get the first child from the `get_children` generator.
297	skip_to_last_child = None # save reference to last child.	20 ✔
298	for skip_to_last_child in node.get_children():	20 ✔
299	pass # skip to last	×
300	return skip_to_last_child # postcondition: node, or None.	20 ✔
301
302
303	def end_setter_from_source(	20 ✔
304	source_code: list[str], pred: NodePredicate, only_consumables: bool = False
305	) -> NodeTransform:
306	"""Returns a function that sets ending locations for a node from source.
307
308	The basic technique is to do the following:
309	1. Find the ending locations for the node based on its last child.
310	2. Starting at that point, iterate through characters in the source code
311	up to and including the first index that satisfies pred.
312
313	pred is a function that takes a string and index and returns a bool,
314	e.g. _is_close_paren
315
316	If only_consumables is True, the search halts when it reaches a non-consumable
317	character that fails pred on the first line.
318	TODO: really the behaviour should be the same for all lines searched for.
319	"""
320
321	def set_endings_from_source(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
322	# Tuple nodes have an end_col_offset that includes the end paren,
323	# but their col_offset does not include the start paren.
324	# To address this, we override the Tuple node's end_col_offset.
325	if isinstance(node, nodes.Tuple):	20 ✔
326	set_from_last_child(node)	20 ✔
327
328	# Initialize counters. Note: we need to offset lineno,
329	# since it's 1-indexed.
330	end_col_offset, lineno = node.end_col_offset, node.end_lineno - 1	20 ✔
331
332	# First, search the remaining part of the current end line.
333	for j in range(end_col_offset, len(source_code[lineno])):	20 ✔
334	if source_code[lineno][j] == "#":	20 ✔
335	break # skip over comment lines	20 ✔
336	if pred(source_code[lineno], j, node):	20 ✔
337	node.end_col_offset = j + 1	20 ✔
338	return node	20 ✔
339	elif only_consumables and source_code[lineno][j] not in CONSUMABLES:	20 ✔
340	return node	20 ✔
341
342	# If that doesn't work, search remaining lines
343	for i in range(lineno + 1, len(source_code)):	20 ✔
344	# Search each character
345	for j in range(len(source_code[i])):	20 ✔
346	if source_code[i][j] == "#":	20 ✔
347	break # skip over comment lines	20 ✔
348	if pred(source_code[i], j, node):	20 ✔
349	node.end_col_offset, node.end_lineno = j + 1, i + 1	20 ✔
350	return node	20 ✔
351	# only consume inert characters.
352	elif source_code[i][j] not in CONSUMABLES:	20 ✔
353	return node	20 ✔
354	return node	20 ✔
355
356	return set_endings_from_source	20 ✔
357
358
359	def start_setter_from_source(source_code: list[str], pred: NodePredicate) -> NodeTransform:	20 ✔
360	"""Returns a function that sets start locations for a node from source.
361	Recall `source_code`, `pred` are within the lexical scope of the returned function.
362
363	The basic technique is to do the following:
364	1. Find the start locations for the node (already set).
365	2. Starting at that point, iterate through characters in the source code
366	in reverse until reaching the first index that satisfies pred.
367
368	pred is a function that takes a string and index and returns a bool,
369	e.g. _is_open_paren
370	"""
371
372	def set_start_from_source(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
373	# Initialize counters. Note: fromlineno is 1-indexed.
374	col_offset, lineno = node.col_offset, node.fromlineno - 1	20 ✔
375
376	# First, search the remaining part of the current start line
377	for j in range(min(len(source_code[lineno]) - 1, col_offset), -1, -1):	20 ✔
378	if pred(source_code[lineno], j, node):	20 ✔
379	node.col_offset = j	20 ✔
380	return node	20 ✔
381
382	# If that doesn't work, search remaining lines
383	for i in range(lineno - 1, -1, -1):	×
384	# Search each character, right-to-left
385	for j in range(len(source_code[i]) - 1, -1, -1):	×
386	if pred(source_code[i], j, node):	×
387	node.end_col_offset, node.end_lineno = j, i + 1	×
388	return node	×
389	# only consume inert characters.
390	elif source_code[i][j] not in CONSUMABLES:	×
391	return node	×
392	return node	×
393
394	return set_start_from_source	20 ✔
395
396
397	def add_parens(source_code: list[str]) -> NodeTransform:	20 ✔
398	def h(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
399	_add_parens(source_code)(node)	20 ✔
400
401	return h	20 ✔
402
403
404	def _add_parens(source_code: list[str]) -> NodeTransform:	20 ✔
405	def h(node: nodes.NodeNG) -> nodes.NodeNG:	20 ✔
406	# Initialize counters. Note: fromlineno is 1-indexed.
407	prev = node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset	20 ✔
408	while True:	20 ✔
409	col_offset, lineno = node.col_offset, node.fromlineno - 1	20 ✔
410	end_col_offset, end_lineno = node.end_col_offset, node.end_lineno - 1	20 ✔
411
412	# First, search the remaining part of the current start line
413	prev_char, new_lineno, new_coloffset = None, None, None	20 ✔
414	for j in range(col_offset - 1, -1, -1):	20 ✔
415	if source_code[lineno][j] in CONSUMABLES or source_code[lineno][j] == ",":	20 ✔
416	continue	20 ✔
417	else:
418	prev_char, new_lineno, new_coloffset = source_code[lineno][j], lineno, j	20 ✔
419	break	20 ✔
420
421	if prev_char is None:	20 ✔
422	# Search remaining lines
423	for i in range(lineno - 1, -1, -1):	20 ✔
424	# Search each character, right-to-left
425	for j in range(len(source_code[i]) - 1, -1, -1):	20 ✔
426	if source_code[i][j] in CONSUMABLES or source_code[i][j] == ",":	20 ✔
427	continue	20 ✔
428	else:
429	prev_char, new_lineno, new_coloffset = source_code[i][j], i, j	20 ✔
430
431	break	20 ✔
432	if prev_char is not None:	20 ✔
433	break	20 ✔
434
435	if prev_char is None or new_lineno is None or new_coloffset is None:	20 ✔
436	break	20 ✔
437
438	if prev_char != "(":	20 ✔
439	# No enclosing parentheses
440	break	20 ✔
441
442	# Now search for matching ')'
443	next_char, new_end_lineno, new_end_coloffset = None, None, None	20 ✔
444	for j in range(end_col_offset, len(source_code[end_lineno])):	20 ✔
445	if source_code[end_lineno][j] == "#":	20 ✔
446	break # skip over comment lines	×
447	elif source_code[end_lineno][j] in CONSUMABLES:	20 ✔
448	continue	20 ✔
449	else:
450	next_char, new_end_lineno, new_end_coloffset = (	20 ✔
451	source_code[end_lineno][j],
452	end_lineno,
453	j,
454	)
455	break	20 ✔
456
457	if next_char is None:	20 ✔
458	# Search remaining lines
459	for i in range(end_lineno + 1, len(source_code)):	20 ✔
460	# Search each character
461	for j in range(len(source_code[i])):	20 ✔
462	if source_code[i][j] == "#":	20 ✔
463	break # skip over comment lines	×
464	elif source_code[i][j] in CONSUMABLES:	20 ✔
465	continue	20 ✔
466	else:
467	next_char, new_end_lineno, new_end_coloffset = source_code[i][j], i, j	20 ✔
468	break	20 ✔
469	if next_char is not None:	20 ✔
470	break	20 ✔
471
472	if next_char is None or new_end_lineno is None or new_end_coloffset is None:	20 ✔
NEW 473	break	×
474
475	if next_char != ")":	20 ✔
476	break	20 ✔
477
478	# At this point, an enclosing pair of parentheses has been found
479	prev = node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset	20 ✔
480	node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset = (	20 ✔
481	new_lineno + 1,
482	new_coloffset,
483	new_end_lineno + 1,
484	new_end_coloffset + 1,
485	)
486
487	# Go back by 1 set of parentheses if inside a function call.
488	if isinstance(node.parent, nodes.Call) and len(node.parent.args) == 1:	20 ✔
489	node.fromlineno, node.col_offset, node.end_lineno, node.end_col_offset = prev	20 ✔
490
491	return node	20 ✔
492
493	return h	20 ✔
494
495
496	# Make this module a pylint plugin
497	def register(linter: PyLinter) -> None:	20 ✔
498	"""Patch linter to apply message transform with source code."""
499	old_get_ast = linter.get_ast	20 ✔
500
501	def new_get_ast(filepath: str, modname: str, data: str) -> Any:	20 ✔
502	ast = old_get_ast(filepath, modname, data)	20 ✔
503	if ast is not None:	20 ✔
504	with open(filepath, encoding="utf-8") as f:	20 ✔
505	source_code = f.readlines()	20 ✔
506	ending_transformer = init_register_ending_setters(source_code)	20 ✔
507	ending_transformer.visit(ast)	20 ✔
508	return ast	20 ✔
509
510	linter.get_ast = new_get_ast # type: ignore[assignment, method-assign]	20 ✔

pyta-uoft / pyta / 36646375051

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous