Fix get_stack_info for python3.12 (#421)

tripy/tests/frontend/test_tensor.py

@@ -127,7 +127,7 @@ def test_stack_info_is_populated(self, build_func, expected_line_number):
             function=tp.Tensor.raw_init.__name__,
             code=None,
             _dispatch_target="",
-            column_range=None,
+            column_range=(25, 30) if sys.version_info >= (3, 11) else None,
         )
         assert a.stack_info[a.stack_info.get_first_user_frame_index()] == SourceInfo(
             __name__,
@@ -136,7 +136,7 @@ def test_stack_info_is_populated(self, build_func, expected_line_number):
             function=build_func.__name__,
             code=inspect.getsource(build_func).rstrip(),
             _dispatch_target="",
-            column_range=None,
+            column_range=(0, 0) if sys.version_info >= (3, 11) else None,
         )
 
     def test_eval_of_storage_tensor_is_nop(self):

tripy/tripy/utils/stack_info.py

@@ -119,14 +119,16 @@ def get_stack_info(include_code_index: int = None) -> StackInfo:
             source_info._dispatch_target = frame.f_locals["key"]
 
         try:
-            # In Python 3.11, frames contain column offset information.
             frame.f_code.co_positions
         except AttributeError:
             pass
         else:
-            # There are twice as many instructions as co_positions()
-            _, _, start, end = frame.f_code.co_positions()[frame.f_lasti // 2]
-            source_info.column_range = (start, end)
+            # For 3.11+, positions is an iterator of (line, end_line, column, end_column)
+            for pos in frame.f_code.co_positions():
+                if pos[0] == frame.f_lineno:
+                    _, _, start, end = pos
+                    source_info.column_range = (start, end)
+                    break
 
         stack_info.append(source_info)
         frame = frame.f_back