Copy more code attributes

ionelmc · Mar 8, 2020 · 43ed15d · 43ed15d
1 parent 0464b02
commit 43ed15d
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 9 deletions.
diff --git a/AUTHORS.rst b/AUTHORS.rst
@@ -10,3 +10,4 @@ Authors
 * Elliott Sales de Andrade - https://github.com/QuLogic
 * Victor Stinner - https://github.com/vstinner
 * Guido Imperiale - https://github.com/crusaderky
+* Ivanq - https://github.com/imachug
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -2,6 +2,12 @@
 Changelog
 =========
 
+1.7.0 (2020-03-07)
+~~~~~~~~~~~~~~~~~~
+
+* Add more attributes to ``Frame`` and ``Code`` objects for pytest compatibility. Contributed by Ivanq in
+  `#58 <https://github.com/ionelmc/python-tblib/pull/58>`_.
+
 1.6.0 (2019-12-07)
 ~~~~~~~~~~~~~~~~~~
 

diff --git a/README.rst b/README.rst
@@ -440,19 +440,23 @@ json.JSONDecoder::
     ...     pprint(tb_dict)
     {'tb_frame': {'f_code': {'co_filename': '<doctest README.rst[...]>',
                              'co_name': '<module>'},
-                  'f_globals': {'__name__': '__main__'}},
+                  'f_globals': {'__name__': '__main__'},
+                  'f_lineno': 5},
      'tb_lineno': 2,
-     'tb_next': {'tb_frame': {'f_code': {'co_filename': ...
+     'tb_next': {'tb_frame': {'f_code': {'co_filename': ...,
                                          'co_name': 'inner_2'},
-                              'f_globals': {'__name__': '__main__'}},
+                              'f_globals': {'__name__': '__main__'},
+                              'f_lineno': 2},
                  'tb_lineno': 2,
-                 'tb_next': {'tb_frame': {'f_code': {'co_filename': ...
+                 'tb_next': {'tb_frame': {'f_code': {'co_filename': ...,
                                                      'co_name': 'inner_1'},
-                                          'f_globals': {'__name__': '__main__'}},
+                                          'f_globals': {'__name__': '__main__'},
+                                          'f_lineno': 2},
                              'tb_lineno': 2,
-                             'tb_next': {'tb_frame': {'f_code': {'co_filename': ...
+                             'tb_next': {'tb_frame': {'f_code': {'co_filename': ...,
                                                                  'co_name': 'inner_0'},
-                                                      'f_globals': {'__name__': '__main__'}},
+                                                      'f_globals': {'__name__': '__main__'},
+                                                      'f_lineno': 2},
                                          'tb_lineno': 2,
                                          'tb_next': None}}}}
 

diff --git a/src/tblib/__init__.py b/src/tblib/__init__.py
@@ -43,16 +43,25 @@ class Code(object):
     def __init__(self, code):
         self.co_filename = code.co_filename
         self.co_name = code.co_name
+        self.co_argcount = 0
+        self.co_kwonlyargcount = 0
+        self.co_varnames = ()
+        self.co_nlocals = 0
+        self.co_stacksize = 0
+        self.co_flags = 64
+        self.co_firstlineno = 0
 
 
 class Frame(object):
     def __init__(self, frame):
+        self.f_locals = {}
         self.f_globals = {
             k: v
             for k, v in frame.f_globals.items()
             if k in ("__file__", "__name__")
         }
         self.f_code = Code(frame.f_code)
+        self.f_lineno = frame.f_lineno
 
     def clear(self):
         # For compatibility with PyPy 3.5;
@@ -161,6 +170,7 @@ def as_dict(self):
         frame = {
             'f_globals': self.tb_frame.f_globals,
             'f_code': code,
+            'f_lineno': self.tb_frame.f_lineno,
         }
         return {
             'tb_frame': frame,
@@ -183,6 +193,7 @@ def from_dict(cls, dct):
         frame = _AttrDict(
             f_globals=dct['tb_frame']['f_globals'],
             f_code=code,
+            f_lineno=dct['tb_frame']['f_lineno'],
         )
         tb = _AttrDict(
             tb_frame=frame,
@@ -222,6 +233,7 @@ def from_string(cls, string, strict=True):
                             __name__='?',
                         ),
                         f_code=_AttrDict(frame),
+                        f_lineno=int(frame['tb_lineno']),
                     ),
                     tb_next=previous,
                 )

diff --git a/tests/test_tblib.py b/tests/test_tblib.py
@@ -33,18 +33,21 @@ def test_parse_traceback():
         "tb_frame": {
             "f_code": {"co_filename": "file1", "co_name": "<module>"},
             "f_globals": {"__file__": "file1", "__name__": "?"},
+            "f_lineno": 123,
         },
         "tb_lineno": 123,
         "tb_next": {
             "tb_frame": {
                 "f_code": {"co_filename": "file2", "co_name": "???"},
                 "f_globals": {"__file__": "file2", "__name__": "?"},
+                "f_lineno": 234,
             },
             "tb_lineno": 234,
             "tb_next": {
                 "tb_frame": {
                     "f_code": {"co_filename": "file3", "co_name": "function3"},
                     "f_globals": {"__file__": "file3", "__name__": "?"},
+                    "f_lineno": 345,
                 },
                 "tb_lineno": 345,
                 "tb_next": None,
@@ -105,7 +108,7 @@ def test_raise():
 
     result = testdir.runpytest_subprocess('--tb=short', '-vv', test)
     result.stdout.fnmatch_lines([
-        'test_pytest_integration.py:15: in test_raise',
+        'test_pytest_integration.py:*: in test_raise',
         '    six.reraise(RuntimeError, RuntimeError(), pytb)',
         'file1:123: in <module>',
         '    ???',
@@ -128,7 +131,7 @@ def test_raise():
     result = testdir.runpytest_subprocess('--tb=native', '-vv', test)
     result.stdout.fnmatch_lines([
         'Traceback (most recent call last):',
-        '  File "*test_pytest_integration.py", line 15, in test_raise',
+        '  File "*test_pytest_integration.py", line *, in test_raise',
         '    six.reraise(RuntimeError, RuntimeError(), pytb)',
         '  File "file1", line 123, in <module>',
         '  File "file2", line 234, in ???',