1010from types import CodeType
1111from typing import Iterable , Optional
1212
13- from coverage .types import TArc , TOffset
13+ from coverage .types import TArc , TLineNo , TOffset
1414
1515
1616def code_objects (code : CodeType ) -> Iterable [CodeType ]:
@@ -99,7 +99,7 @@ def walk(
9999TBranchTrails = dict [TOffset , TBranchTrailsOneSource ]
100100
101101
102- def branch_trails (code : CodeType ) -> TBranchTrails :
102+ def branch_trails (code : CodeType , multiline_map : dict [ TLineNo , TLineNo ] ) -> TBranchTrails :
103103 """
104104 Calculate branch trails for `code`.
105105
@@ -128,6 +128,7 @@ def branch_trails(code: CodeType) -> TBranchTrails:
128128 from_line = inst .line_number
129129 if from_line is None :
130130 continue
131+ from_line = multiline_map .get (from_line , from_line )
131132
132133 def add_one_branch_trail (
133134 trails : TBranchTrailsOneSource ,
@@ -138,8 +139,11 @@ def add_one_branch_trail(
138139 to_line = None
139140 for inst2 in iwalker .walk (start_at = start_at , follow_jumps = True ):
140141 inst_offsets .add (inst2 .offset )
141- if inst2 .line_number and inst2 .line_number != from_line :
142- to_line = inst2 .line_number
142+ l2 = inst2 .line_number
143+ if l2 is not None :
144+ l2 = multiline_map .get (l2 , l2 )
145+ if l2 and l2 != from_line :
146+ to_line = l2
143147 break
144148 elif inst2 .jump_target and (inst2 .opcode not in ALWAYS_JUMPS ):
145149 break
0 commit comments