|
2 | 2 | from typing import Any, Collection, Dict, List, Optional, Union, TYPE_CHECKING |
3 | 3 |
|
4 | 4 | if TYPE_CHECKING: |
5 | | - from ..language.ast import Node # noqa: F401 |
| 5 | + from ..language.ast import Node, Location # noqa: F401 |
6 | 6 | from ..language.location import SourceLocation # noqa: F401 |
7 | 7 | from ..language.source import Source # noqa: F401 |
8 | 8 |
|
@@ -95,27 +95,24 @@ def __init__( |
95 | 95 | if nodes and not isinstance(nodes, list): |
96 | 96 | nodes = [nodes] # type: ignore |
97 | 97 | self.nodes = nodes or None # type: ignore |
| 98 | + node_locations = ( |
| 99 | + [node.loc for node in nodes if node.loc] if nodes else [] # type: ignore |
| 100 | + ) |
98 | 101 | self.source = source |
99 | | - if not source and nodes: |
100 | | - node = nodes[0] # type: ignore |
101 | | - if node and node.loc and node.loc.source: |
102 | | - self.source = node.loc.source |
103 | | - if not positions and nodes: |
104 | | - positions = [node.loc.start for node in nodes if node.loc] # type: ignore |
| 102 | + if not source and node_locations: |
| 103 | + loc = node_locations[0] |
| 104 | + if loc.source: # pragma: no cover else |
| 105 | + self.source = loc.source |
| 106 | + if not positions and node_locations: |
| 107 | + positions = [loc.start for loc in node_locations] |
105 | 108 | self.positions = positions or None |
106 | 109 | if positions and source: |
107 | 110 | locations: Optional[List["SourceLocation"]] = [ |
108 | 111 | source.get_location(pos) for pos in positions |
109 | 112 | ] |
110 | | - elif nodes: |
111 | | - locations = [ |
112 | | - node.loc.source.get_location(node.loc.start) |
113 | | - for node in nodes # type: ignore |
114 | | - if node.loc |
115 | | - ] |
116 | 113 | else: |
117 | | - locations = None |
118 | | - self.locations = locations |
| 114 | + locations = [loc.source.get_location(loc.start) for loc in node_locations] |
| 115 | + self.locations = locations or None |
119 | 116 | if path and not isinstance(path, list): |
120 | 117 | path = list(path) |
121 | 118 | self.path = path or None # type: ignore |
|
0 commit comments