Merge pull request #7980 from bluetech/code-changes

code: a few minor improvements

3 files changed, 57 insertions, 31 deletions

diff --git a/src/_pytest/_code/code.py b/src/_pytest/_code/code.py
index 2371b44d9..423069330 100644
--- a/src/_pytest/_code/code.py
+++ b/src/_pytest/_code/code.py
@@ -54,15 +54,14 @@ if TYPE_CHECKING:
 class Code:
     """Wrapper around Python code objects."""
 
-    def __init__(self, rawcode) -> None:
-        if not hasattr(rawcode, "co_filename"):
-            rawcode = getrawcode(rawcode)
-        if not isinstance(rawcode, CodeType):
-            raise TypeError(f"not a code object: {rawcode!r}")
-        self.filename = rawcode.co_filename
-        self.firstlineno = rawcode.co_firstlineno - 1
-        self.name = rawcode.co_name
-        self.raw = rawcode
+    __slots__ = ("raw",)
+
+    def __init__(self, obj: CodeType) -> None:
+        self.raw = obj
+
+    @classmethod
+    def from_function(cls, obj: object) -> "Code":
+        return cls(getrawcode(obj))
 
     def __eq__(self, other):
         return self.raw == other.raw
@@ -71,6 +70,14 @@ class Code:
     __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
         case of ``OSError`` / non-existing file."""
@@ -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
-    exprinfo = None
+    __slots__ = ("_rawentry", "_excinfo", "_repr_style")
 
     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]
diff --git a/src/_pytest/_code/source.py b/src/_pytest/_code/source.py
index c63a42360..6f54057c0 100644
--- a/src/_pytest/_code/source.py
+++ b/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:
-            rawcode = getrawcode(obj)
-            src = inspect.getsource(rawcode)
+            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)
-        obj = getattr(obj, "__code__", obj)
-        if trycall and not hasattr(obj, "co_firstlineno"):
-            if hasattr(obj, "__call__") and not inspect.isclass(obj):
-                x = getrawcode(obj.__call__, trycall=False)
-                if hasattr(x, "co_firstlineno"):
-                    return x
-        return obj
+        pass
+    if trycall:
+        call = getattr(obj, "__call__", None)
+        if call and not isinstance(obj, type):
+            return getrawcode(call, trycall=False)
+    raise TypeError(f"could not get code object for {obj!r}")
 
 
 def deindent(lines: Iterable[str]) -> List[str]:
diff --git a/src/_pytest/python.py b/src/_pytest/python.py
index 9e988032d..a3eaa5823 100644
--- a/src/_pytest/python.py
+++ b/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)