code: simplify Code construction

This commit is contained in:
Ran Benita 2020-10-27 16:07:03 +02:00
parent 6506f016ac
commit 531416cc5a
7 changed files with 40 additions and 45 deletions

View File

@ -56,12 +56,12 @@ class Code:
__slots__ = ("raw",) __slots__ = ("raw",)
def __init__(self, rawcode) -> None: def __init__(self, obj: CodeType) -> None:
if not hasattr(rawcode, "co_filename"): self.raw = obj
rawcode = getrawcode(rawcode)
if not isinstance(rawcode, CodeType): @classmethod
raise TypeError(f"not a code object: {rawcode!r}") def from_function(cls, obj: object) -> "Code":
self.raw = rawcode return cls(getrawcode(obj))
def __eq__(self, other): def __eq__(self, other):
return self.raw == other.raw return self.raw == other.raw
@ -1196,7 +1196,7 @@ def getfslineno(obj: object) -> Tuple[Union[str, py.path.local], int]:
obj = obj.place_as # type: ignore[attr-defined] obj = obj.place_as # type: ignore[attr-defined]
try: try:
code = Code(obj) code = Code.from_function(obj)
except TypeError: except TypeError:
try: try:
fn = inspect.getsourcefile(obj) or inspect.getfile(obj) # type: ignore[arg-type] fn = inspect.getsourcefile(obj) or inspect.getfile(obj) # type: ignore[arg-type]

View File

@ -2,6 +2,7 @@ import ast
import inspect import inspect
import textwrap import textwrap
import tokenize import tokenize
import types
import warnings import warnings
from bisect import bisect_right from bisect import bisect_right
from typing import Iterable from typing import Iterable
@ -29,8 +30,11 @@ class Source:
elif isinstance(obj, str): elif isinstance(obj, str):
self.lines = deindent(obj.split("\n")) self.lines = deindent(obj.split("\n"))
else: else:
try:
rawcode = getrawcode(obj) rawcode = getrawcode(obj)
src = inspect.getsource(rawcode) src = inspect.getsource(rawcode)
except TypeError:
src = inspect.getsource(obj) # type: ignore[arg-type]
self.lines = deindent(src.split("\n")) self.lines = deindent(src.split("\n"))
def __eq__(self, other: object) -> bool: def __eq__(self, other: object) -> bool:
@ -122,19 +126,17 @@ def findsource(obj) -> Tuple[Optional[Source], int]:
return source, lineno return source, lineno
def getrawcode(obj, trycall: bool = True): def getrawcode(obj: object, trycall: bool = True) -> types.CodeType:
"""Return code object for given function.""" """Return code object for given function."""
try: try:
return obj.__code__ return obj.__code__ # type: ignore[attr-defined,no-any-return]
except AttributeError: 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]: def deindent(lines: Iterable[str]) -> List[str]:

View File

@ -1647,7 +1647,7 @@ class Function(PyobjMixin, nodes.Item):
def _prunetraceback(self, excinfo: ExceptionInfo[BaseException]) -> None: def _prunetraceback(self, excinfo: ExceptionInfo[BaseException]) -> None:
if hasattr(self, "_obj") and not self.config.getoption("fulltrace", False): 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 path, firstlineno = code.path, code.firstlineno
traceback = excinfo.traceback traceback = excinfo.traceback
ntraceback = traceback.cut(path=path, firstlineno=firstlineno) ntraceback = traceback.cut(path=path, firstlineno=firstlineno)

View File

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

View File

@ -147,7 +147,7 @@ class TestTraceback_f_g_h:
] ]
def test_traceback_cut(self): def test_traceback_cut(self):
co = _pytest._code.Code(f) co = _pytest._code.Code.from_function(f)
path, firstlineno = co.path, co.firstlineno path, firstlineno = co.path, co.firstlineno
traceback = self.excinfo.traceback traceback = self.excinfo.traceback
newtraceback = traceback.cut(path=path, firstlineno=firstlineno) newtraceback = traceback.cut(path=path, firstlineno=firstlineno)
@ -290,7 +290,7 @@ class TestTraceback_f_g_h:
excinfo = pytest.raises(ValueError, f) excinfo = pytest.raises(ValueError, f)
tb = excinfo.traceback tb = excinfo.traceback
entry = tb.getcrashentry() entry = tb.getcrashentry()
co = _pytest._code.Code(h) co = _pytest._code.Code.from_function(h)
assert entry.frame.code.path == co.path assert entry.frame.code.path == co.path
assert entry.lineno == co.firstlineno + 1 assert entry.lineno == co.firstlineno + 1
assert entry.frame.code.name == "h" assert entry.frame.code.name == "h"
@ -307,7 +307,7 @@ class TestTraceback_f_g_h:
excinfo = pytest.raises(ValueError, f) excinfo = pytest.raises(ValueError, f)
tb = excinfo.traceback tb = excinfo.traceback
entry = tb.getcrashentry() entry = tb.getcrashentry()
co = _pytest._code.Code(g) co = _pytest._code.Code.from_function(g)
assert entry.frame.code.path == co.path assert entry.frame.code.path == co.path
assert entry.lineno == co.firstlineno + 2 assert entry.lineno == co.firstlineno + 2
assert entry.frame.code.name == "g" assert entry.frame.code.name == "g"

View File

@ -16,8 +16,8 @@ import py.path
import pytest import pytest
from _pytest._code import Code from _pytest._code import Code
from _pytest._code import Frame from _pytest._code import Frame
from _pytest._code import Source
from _pytest._code import getfslineno from _pytest._code import getfslineno
from _pytest._code import Source
def test_source_str_function() -> None: def test_source_str_function() -> None:
@ -291,7 +291,7 @@ def test_source_of_class_at_eof_without_newline(tmpdir, _sys_snapshot) -> None:
# does not return the "x = 1" last line. # does not return the "x = 1" last line.
source = Source( source = Source(
""" """
class A(object): class A:
def method(self): def method(self):
x = 1 x = 1
""" """
@ -374,14 +374,6 @@ def test_getfslineno() -> None:
B.__name__ = B.__qualname__ = "B2" B.__name__ = B.__qualname__ = "B2"
assert getfslineno(B)[1] == -1 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: def test_code_of_object_instance_with_call() -> None:
class A: class A:
@ -393,14 +385,14 @@ def test_code_of_object_instance_with_call() -> None:
def __call__(self) -> None: def __call__(self) -> None:
pass pass
code = Code(WithCall()) code = Code.from_function(WithCall())
assert "pass" in str(code.source()) assert "pass" in str(code.source())
class Hello: class Hello:
def __call__(self) -> None: def __call__(self) -> None:
pass pass
pytest.raises(TypeError, lambda: Code(Hello)) pytest.raises(TypeError, lambda: Code.from_function(Hello))
def getstatement(lineno: int, source) -> Source: def getstatement(lineno: int, source) -> Source:

View File

@ -42,7 +42,7 @@ def getmsg(
f, extra_ns: Optional[Mapping[str, object]] = None, *, must_pass: bool = False f, extra_ns: Optional[Mapping[str, object]] = None, *, must_pass: bool = False
) -> Optional[str]: ) -> Optional[str]:
"""Rewrite the assertions in f, run it, and get the failure message.""" """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) mod = rewrite(src)
code = compile(mod, "<test>", "exec") code = compile(mod, "<test>", "exec")
ns: Dict[str, object] = {} ns: Dict[str, object] = {}