Merge pull request #7980 from bluetech/code-changes

code: a few minor improvements

src/_pytest/_code/code.py

@@ -54,15 +54,14 @@ if TYPE_CHECKING:
 class Code:
     """Wrapper around Python code objects."""
 
-    def __init__(self, rawcode) -> None:
+    __slots__ = ("raw",)
-        if not hasattr(rawcode, "co_filename"):
+
-            rawcode = getrawcode(rawcode)
+    def __init__(self, obj: CodeType) -> None:
-        if not isinstance(rawcode, CodeType):
+        self.raw = obj
-            raise TypeError(f"not a code object: {rawcode!r}")
+
-        self.filename = rawcode.co_filename
+    @classmethod
-        self.firstlineno = rawcode.co_firstlineno - 1
+    def from_function(cls, obj: object) -> "Code":
-        self.name = rawcode.co_name
+        return cls(getrawcode(obj))
-        self.raw = rawcode
 
     def __eq__(self, other):
         return self.raw == other.raw
@@ -70,6 +69,14 @@ class Code:
     # Ignore type because of https://github.com/python/mypy/issues/4266.
     __hash__ = None  # type: ignore
 
+    @property
+    def firstlineno(self) -> int:
+        return self.raw.co_firstlineno - 1
+
+    @property
+    def name(self) -> str:
+        return self.raw.co_name
+
     @property
     def path(self) -> Union[py.path.local, str]:
         """Return a path object pointing to source code, or an ``str`` in
@@ -117,12 +124,26 @@ class Frame:
     """Wrapper around a Python frame holding f_locals and f_globals
     in which expressions can be evaluated."""
 
+    __slots__ = ("raw",)
+
     def __init__(self, frame: FrameType) -> None:
-        self.lineno = frame.f_lineno - 1
-        self.f_globals = frame.f_globals
-        self.f_locals = frame.f_locals
         self.raw = frame
-        self.code = Code(frame.f_code)
+
+    @property
+    def lineno(self) -> int:
+        return self.raw.f_lineno - 1
+
+    @property
+    def f_globals(self) -> Dict[str, Any]:
+        return self.raw.f_globals
+
+    @property
+    def f_locals(self) -> Dict[str, Any]:
+        return self.raw.f_locals
+
+    @property
+    def code(self) -> Code:
+        return Code(self.raw.f_code)
 
     @property
     def statement(self) -> "Source":
@@ -164,17 +185,20 @@ class Frame:
 class TracebackEntry:
     """A single entry in a Traceback."""
 
-    _repr_style: Optional['Literal["short", "long"]'] = None
+    __slots__ = ("_rawentry", "_excinfo", "_repr_style")
-    exprinfo = None
 
     def __init__(
         self,
         rawentry: TracebackType,
         excinfo: Optional["ReferenceType[ExceptionInfo[BaseException]]"] = None,
     ) -> None:
-        self._excinfo = excinfo
         self._rawentry = rawentry
-        self.lineno = rawentry.tb_lineno - 1
+        self._excinfo = excinfo
+        self._repr_style: Optional['Literal["short", "long"]'] = None
+
+    @property
+    def lineno(self) -> int:
+        return self._rawentry.tb_lineno - 1
 
     def set_repr_style(self, mode: "Literal['short', 'long']") -> None:
         assert mode in ("short", "long")
@@ -1172,7 +1196,7 @@ def getfslineno(obj: object) -> Tuple[Union[str, py.path.local], int]:
         obj = obj.place_as  # type: ignore[attr-defined]
 
     try:
-        code = Code(obj)
+        code = Code.from_function(obj)
     except TypeError:
         try:
             fn = inspect.getsourcefile(obj) or inspect.getfile(obj)  # type: ignore[arg-type]

src/_pytest/_code/source.py

@@ -2,6 +2,7 @@ import ast
 import inspect
 import textwrap
 import tokenize
+import types
 import warnings
 from bisect import bisect_right
 from typing import Iterable
@@ -29,8 +30,11 @@ class Source:
         elif isinstance(obj, str):
             self.lines = deindent(obj.split("\n"))
         else:
+            try:
                 rawcode = getrawcode(obj)
                 src = inspect.getsource(rawcode)
+            except TypeError:
+                src = inspect.getsource(obj)  # type: ignore[arg-type]
             self.lines = deindent(src.split("\n"))
 
     def __eq__(self, other: object) -> bool:
@@ -122,19 +126,17 @@ def findsource(obj) -> Tuple[Optional[Source], int]:
     return source, lineno
 
 
-def getrawcode(obj, trycall: bool = True):
+def getrawcode(obj: object, trycall: bool = True) -> types.CodeType:
     """Return code object for given function."""
     try:
-        return obj.__code__
+        return obj.__code__  # type: ignore[attr-defined,no-any-return]
     except AttributeError:
-        obj = getattr(obj, "f_code", obj)
+        pass
-        obj = getattr(obj, "__code__", obj)
+    if trycall:
-        if trycall and not hasattr(obj, "co_firstlineno"):
+        call = getattr(obj, "__call__", None)
-            if hasattr(obj, "__call__") and not inspect.isclass(obj):
+        if call and not isinstance(obj, type):
-                x = getrawcode(obj.__call__, trycall=False)
+            return getrawcode(call, trycall=False)
-                if hasattr(x, "co_firstlineno"):
+    raise TypeError(f"could not get code object for {obj!r}")
-                    return x
-        return obj
 
 
 def deindent(lines: Iterable[str]) -> List[str]:

src/_pytest/python.py

@@ -1648,7 +1648,7 @@ class Function(PyobjMixin, nodes.Item):
 
     def _prunetraceback(self, excinfo: ExceptionInfo[BaseException]) -> None:
         if hasattr(self, "_obj") and not self.config.getoption("fulltrace", False):
-            code = _pytest._code.Code(get_real_func(self.obj))
+            code = _pytest._code.Code.from_function(get_real_func(self.obj))
             path, firstlineno = code.path, code.firstlineno
             traceback = excinfo.traceback
             ntraceback = traceback.cut(path=path, firstlineno=firstlineno)

testing/code/test_code.py

@@ -28,11 +28,12 @@ def test_code_gives_back_name_for_not_existing_file() -> None:
     assert code.fullsource is None
 
 
-def test_code_with_class() -> None:
+def test_code_from_function_with_class() -> None:
     class A:
         pass
 
-    pytest.raises(TypeError, Code, A)
+    with pytest.raises(TypeError):
+        Code.from_function(A)
 
 
 def x() -> None:
@@ -40,13 +41,13 @@ def x() -> None:
 
 
 def test_code_fullsource() -> None:
-    code = Code(x)
+    code = Code.from_function(x)
     full = code.fullsource
     assert "test_code_fullsource()" in str(full)
 
 
 def test_code_source() -> None:
-    code = Code(x)
+    code = Code.from_function(x)
     src = code.source()
     expected = """def x() -> None:
     raise NotImplementedError()"""
@@ -73,7 +74,7 @@ def test_getstatement_empty_fullsource() -> None:
 
 
 def test_code_from_func() -> None:
-    co = Code(test_frame_getsourcelineno_myself)
+    co = Code.from_function(test_frame_getsourcelineno_myself)
     assert co.firstlineno
     assert co.path
 
@@ -92,25 +93,25 @@ def test_code_getargs() -> None:
     def f1(x):
         raise NotImplementedError()
 
-    c1 = Code(f1)
+    c1 = Code.from_function(f1)
     assert c1.getargs(var=True) == ("x",)
 
     def f2(x, *y):
         raise NotImplementedError()
 
-    c2 = Code(f2)
+    c2 = Code.from_function(f2)
     assert c2.getargs(var=True) == ("x", "y")
 
     def f3(x, **z):
         raise NotImplementedError()
 
-    c3 = Code(f3)
+    c3 = Code.from_function(f3)
     assert c3.getargs(var=True) == ("x", "z")
 
     def f4(x, *y, **z):
         raise NotImplementedError()
 
-    c4 = Code(f4)
+    c4 = Code.from_function(f4)
     assert c4.getargs(var=True) == ("x", "y", "z")
 
 

testing/code/test_excinfo.py

@@ -147,7 +147,7 @@ class TestTraceback_f_g_h:
         ]
 
     def test_traceback_cut(self):
-        co = _pytest._code.Code(f)
+        co = _pytest._code.Code.from_function(f)
         path, firstlineno = co.path, co.firstlineno
         traceback = self.excinfo.traceback
         newtraceback = traceback.cut(path=path, firstlineno=firstlineno)
@@ -290,7 +290,7 @@ class TestTraceback_f_g_h:
         excinfo = pytest.raises(ValueError, f)
         tb = excinfo.traceback
         entry = tb.getcrashentry()
-        co = _pytest._code.Code(h)
+        co = _pytest._code.Code.from_function(h)
         assert entry.frame.code.path == co.path
         assert entry.lineno == co.firstlineno + 1
         assert entry.frame.code.name == "h"
@@ -307,7 +307,7 @@ class TestTraceback_f_g_h:
         excinfo = pytest.raises(ValueError, f)
         tb = excinfo.traceback
         entry = tb.getcrashentry()
-        co = _pytest._code.Code(g)
+        co = _pytest._code.Code.from_function(g)
         assert entry.frame.code.path == co.path
         assert entry.lineno == co.firstlineno + 2
         assert entry.frame.code.name == "g"
@@ -747,7 +747,6 @@ raise ValueError()
         from _pytest._code.code import Code
 
         monkeypatch.setattr(Code, "path", "bogus")
-        excinfo.traceback[0].frame.code.path = "bogus"  # type: ignore[misc]
         p = FormattedExcinfo(style="short")
         reprtb = p.repr_traceback_entry(excinfo.traceback[-2])
         lines = reprtb.lines

testing/code/test_source.py

@@ -13,8 +13,9 @@ from typing import Optional
 
 import py.path
 
-import _pytest._code
 import pytest
+from _pytest._code import Code
+from _pytest._code import Frame
 from _pytest._code import getfslineno
 from _pytest._code import Source
 
@@ -35,7 +36,7 @@ def test_source_str_function() -> None:
 
 
 def test_source_from_function() -> None:
-    source = _pytest._code.Source(test_source_str_function)
+    source = Source(test_source_str_function)
     assert str(source).startswith("def test_source_str_function() -> None:")
 
 
@@ -44,13 +45,13 @@ def test_source_from_method() -> None:
         def test_method(self):
             pass
 
-    source = _pytest._code.Source(TestClass().test_method)
+    source = Source(TestClass().test_method)
     assert source.lines == ["def test_method(self):", "    pass"]
 
 
 def test_source_from_lines() -> None:
     lines = ["a \n", "b\n", "c"]
-    source = _pytest._code.Source(lines)
+    source = Source(lines)
     assert source.lines == ["a ", "b", "c"]
 
 
@@ -58,7 +59,7 @@ def test_source_from_inner_function() -> None:
     def f():
         raise NotImplementedError()
 
-    source = _pytest._code.Source(f)
+    source = Source(f)
     assert str(source).startswith("def f():")
 
 
@@ -220,7 +221,7 @@ def test_getstartingblock_singleline() -> None:
     class A:
         def __init__(self, *args) -> None:
             frame = sys._getframe(1)
-            self.source = _pytest._code.Frame(frame).statement
+            self.source = Frame(frame).statement
 
     x = A("x", "y")
 
@@ -250,8 +251,8 @@ def test_getfuncsource_dynamic() -> None:
     def g():
         pass  # pragma: no cover
 
-    f_source = _pytest._code.Source(f)
+    f_source = Source(f)
-    g_source = _pytest._code.Source(g)
+    g_source = Source(g)
     assert str(f_source).strip() == "def f():\n    raise NotImplementedError()"
     assert str(g_source).strip() == "def g():\n    pass  # pragma: no cover"
 
@@ -268,7 +269,7 @@ def test_getfuncsource_with_multine_string() -> None:
     pass
 """
 '''
-    assert str(_pytest._code.Source(f)) == expected.rstrip()
+    assert str(Source(f)) == expected.rstrip()
 
 
 def test_deindent() -> None:
@@ -288,16 +289,16 @@ def test_deindent() -> None:
 def test_source_of_class_at_eof_without_newline(tmpdir, _sys_snapshot) -> None:
     # this test fails because the implicit inspect.getsource(A) below
     # does not return the "x = 1" last line.
-    source = _pytest._code.Source(
+    source = Source(
         """
-        class A(object):
+        class A:
             def method(self):
                 x = 1
     """
     )
     path = tmpdir.join("a.py")
     path.write(source)
-    s2 = _pytest._code.Source(tmpdir.join("a.py").pyimport().A)
+    s2 = Source(tmpdir.join("a.py").pyimport().A)
     assert str(source).strip() == str(s2).strip()
 
 
@@ -373,39 +374,31 @@ def test_getfslineno() -> None:
     B.__name__ = B.__qualname__ = "B2"
     assert getfslineno(B)[1] == -1
 
-    co = compile("...", "", "eval")
-    assert co.co_filename == ""
-
-    if hasattr(sys, "pypy_version_info"):
-        assert getfslineno(co) == ("", -1)
-    else:
-        assert getfslineno(co) == ("", 0)
-
 
 def test_code_of_object_instance_with_call() -> None:
     class A:
         pass
 
-    pytest.raises(TypeError, lambda: _pytest._code.Source(A()))
+    pytest.raises(TypeError, lambda: Source(A()))
 
     class WithCall:
         def __call__(self) -> None:
             pass
 
-    code = _pytest._code.Code(WithCall())
+    code = Code.from_function(WithCall())
     assert "pass" in str(code.source())
 
     class Hello:
         def __call__(self) -> None:
             pass
 
-    pytest.raises(TypeError, lambda: _pytest._code.Code(Hello))
+    pytest.raises(TypeError, lambda: Code.from_function(Hello))
 
 
 def getstatement(lineno: int, source) -> Source:
     from _pytest._code.source import getstatementrange_ast
 
-    src = _pytest._code.Source(source)
+    src = Source(source)
     ast, start, end = getstatementrange_ast(lineno, src)
     return src[start:end]
 
@@ -637,7 +630,7 @@ def test_getstartingblock_multiline() -> None:
     class A:
         def __init__(self, *args):
             frame = sys._getframe(1)
-            self.source = _pytest._code.Frame(frame).statement
+            self.source = Frame(frame).statement
 
     # fmt: off
     x = A('x',

testing/test_assertrewrite.py

@@ -42,7 +42,7 @@ def getmsg(
     f, extra_ns: Optional[Mapping[str, object]] = None, *, must_pass: bool = False
 ) -> Optional[str]:
     """Rewrite the assertions in f, run it, and get the failure message."""
-    src = "\n".join(_pytest._code.Code(f).source().lines)
+    src = "\n".join(_pytest._code.Code.from_function(f).source().lines)
     mod = rewrite(src)
     code = compile(mod, "<test>", "exec")
     ns: Dict[str, object] = {}