from __future__ import annotations

import types

from abc import abstractmethod

from typing import (

Any,

Dict,

Generic,

Iterator,

List,

Optional,

Sequence,

SupportsIndex,

TypeVar,

Union,

overload,

)

# alias to keep the 'bytecode' variable free

import bytecode as _bytecode

from bytecode.flags import CompilerFlags, infer_flags

from bytecode.instr import (

_UNSET,

UNSET,

BaseInstr,

Instr,

Label,

SetLineno,

TryBegin,

TryEnd,

)

class BaseBytecode:

def __init__(self) -> None:

self.argcount = 0

self.posonlyargcount = 0

self.kwonlyargcount = 0

self.first_lineno = 1

self.name = "<module>"

self.qualname = self.name

self.filename = "<string>"

self.docstring: Union[str, None, _UNSET] = UNSET

# We cannot recreate cellvars/freevars from instructions because of super()

# special-case, which involves an implicit __class__ cell/free variable

# We could try to detect it.

# CPython itself breaks if one aliases super so we could maybe make it work

# but it will require careful design and will be done later in the future.

self.cellvars: List[str] = []

self.freevars: List[str] = []

self._flags: CompilerFlags = CompilerFlags(0)

def _copy_attr_from(self, bytecode: BaseBytecode) -> None:

self.argcount = bytecode.argcount

self.posonlyargcount = bytecode.posonlyargcount

self.kwonlyargcount = bytecode.kwonlyargcount

self.flags = bytecode.flags

self.first_lineno = bytecode.first_lineno

self.name = bytecode.name

self.qualname = bytecode.qualname

self.filename = bytecode.filename

self.docstring = bytecode.docstring

self.cellvars = list(bytecode.cellvars)

self.freevars = list(bytecode.freevars)

def __eq__(self, other: Any) -> bool:

if type(self) is not type(other):

return False

if self.argcount != other.argcount:

return False

if self.posonlyargcount != other.posonlyargcount:

return False

if self.kwonlyargcount != other.kwonlyargcount:

return False

if self.flags != other.flags:

return False

if self.first_lineno != other.first_lineno:

return False

if self.filename != other.filename:

return False

if self.name != other.name:

return False

if self.qualname != other.qualname:

return False

if self.docstring != other.docstring:

return False

if self.cellvars != other.cellvars:

return False

if self.freevars != other.freevars:

return False

if self.compute_stacksize() != other.compute_stacksize():

return False

return True

@property

def flags(self) -> CompilerFlags:

return self._flags

@flags.setter

def flags(self, value: CompilerFlags) -> None:

if not isinstance(value, CompilerFlags):

value = CompilerFlags(value)

self._flags = value

def update_flags(self, *, is_async: Optional[bool] = None) -> None:

# infer_flags reasonably only accept concrete subclasses

self.flags = infer_flags(self, is_async) # type: ignore

@abstractmethod

def compute_stacksize(self, *, check_pre_and_post: bool = True) -> int:

raise NotImplementedError

T = TypeVar("T", bound="_BaseBytecodeList")

U = TypeVar("U")

class _BaseBytecodeList(BaseBytecode, list, Generic[U]):

"""List subclass providing type stable slicing and copying."""

@overload

def __getitem__(self, index: SupportsIndex) -> U: ...

@overload

def __getitem__(self: T, index: slice) -> T: ...

def __getitem__(self, index):

value = super().__getitem__(index)

if isinstance(index, slice):

value = type(self)(value)

value._copy_attr_from(self)

return value

def copy(self: T) -> T:

# This is a list subclass and works

new = type(self)(super().copy()) # type: ignore

new._copy_attr_from(self)

return new

def legalize(self) -> None:

"""Check that all the element of the list are valid and remove SetLineno."""

lineno_pos = []

set_lineno = None

current_lineno = self.first_lineno

for pos, instr in enumerate(self):

if isinstance(instr, SetLineno):

set_lineno = instr.lineno

lineno_pos.append(pos)

continue

# Filter out other pseudo instructions

if not isinstance(instr, BaseInstr):

continue

if set_lineno is not None:

instr.lineno = set_lineno

elif instr.lineno is UNSET:

instr.lineno = current_lineno

elif instr.lineno is not None:

current_lineno = instr.lineno

for i in reversed(lineno_pos):

del self[i]

def _check_instr(self, instr):

raise NotImplementedError()

V = TypeVar("V")

class _InstrList(List[V]):

# Providing a stricter typing for this helper whose use is limited to the __eq__

# implementation is more effort than it is worth.

def _flat(self) -> List:

instructions: List = []

labels = {}

jumps = []

try_begins: Dict[TryBegin, int] = {}

try_jumps = []

offset = 0

instr: Any

for index, instr in enumerate(self):

if isinstance(instr, Label):

instructions.append("label_instr%s" % index)

labels[instr] = offset

elif isinstance(instr, TryBegin):

try_begins.setdefault(instr, len(try_begins))

assert isinstance(instr.target, Label)

try_jumps.append((instr.target, len(instructions)))

instructions.append(instr)

elif isinstance(instr, TryEnd):

instructions.append(("TryEnd", try_begins[instr.entry]))

else:

if isinstance(instr, Instr) and isinstance(instr.arg, Label):

target_label = instr.arg

instr = _bytecode.ConcreteInstr(

instr.name, 0, location=instr.location

)

jumps.append((target_label, instr))

instructions.append(instr)

offset += 1

for target_label, instr in jumps:

instr.arg = labels[target_label]

for target_label, index in try_jumps:

instr = instructions[index]

assert isinstance(instr, TryBegin)

instructions[index] = (

"TryBegin",

try_begins[instr],

labels[target_label],

instr.push_lasti,

)

return instructions

def __eq__(self, other: Any) -> bool:

if not isinstance(other, _InstrList):

other = _InstrList(other)

return self._flat() == other._flat()

class Bytecode(

_InstrList[Union[Instr, Label, TryBegin, TryEnd, SetLineno]],

_BaseBytecodeList[Union[Instr, Label, TryBegin, TryEnd, SetLineno]],

):

def __init__(

self,

instructions: Sequence[Union[Instr, Label, TryBegin, TryEnd, SetLineno]] = (),

) -> None:

BaseBytecode.__init__(self)

self.argnames: List[str] = []

for instr in instructions:

self._check_instr(instr)

self.extend(instructions)

def __iter__(self) -> Iterator[Union[Instr, Label, TryBegin, TryEnd, SetLineno]]:

instructions = super().__iter__()

seen_try_begin = False

for instr in instructions:

self._check_instr(instr)

if isinstance(instr, TryBegin):

if seen_try_begin:

raise RuntimeError("TryBegin pseudo instructions cannot be nested.")

seen_try_begin = True

elif isinstance(instr, TryEnd):

seen_try_begin = False

yield instr

def _check_instr(self, instr: Any) -> None:

if not isinstance(instr, (Label, SetLineno, Instr, TryBegin, TryEnd)):

raise ValueError(

"Bytecode must only contain Label, "

"SetLineno, and Instr objects, "

"but %s was found" % type(instr).__name__

)

def _copy_attr_from(self, bytecode: BaseBytecode) -> None:

super()._copy_attr_from(bytecode)

if isinstance(bytecode, Bytecode):

self.argnames = bytecode.argnames

@staticmethod

def from_code(

code: types.CodeType,

prune_caches: bool = True,

conserve_exception_block_stackdepth: bool = False,

) -> Bytecode:

concrete = _bytecode.ConcreteBytecode.from_code(code)

return concrete.to_bytecode(

prune_caches=prune_caches,

conserve_exception_block_stackdepth=conserve_exception_block_stackdepth,

)

def compute_stacksize(self, *, check_pre_and_post: bool = True) -> int:

cfg = _bytecode.ControlFlowGraph.from_bytecode(self)

return cfg.compute_stacksize(check_pre_and_post=check_pre_and_post)

def to_code(

self,

compute_jumps_passes: Optional[int] = None,

stacksize: Optional[int] = None,

*,

check_pre_and_post: bool = True,

compute_exception_stack_depths: bool = True,

) -> types.CodeType:

# Prevent reconverting the concrete bytecode to bytecode and cfg to do the

# calculation if we need to do it.

if stacksize is None or compute_exception_stack_depths:

cfg = _bytecode.ControlFlowGraph.from_bytecode(self)

stacksize = cfg.compute_stacksize(

check_pre_and_post=check_pre_and_post,

compute_exception_stack_depths=compute_exception_stack_depths,

)

self = cfg.to_bytecode()

compute_exception_stack_depths = False # avoid redoing everything

bc = self.to_concrete_bytecode(

compute_jumps_passes=compute_jumps_passes,

compute_exception_stack_depths=compute_exception_stack_depths,

)

return bc.to_code(

stacksize=stacksize,

compute_exception_stack_depths=compute_exception_stack_depths,

)

def to_concrete_bytecode(

self,

compute_jumps_passes: Optional[int] = None,

compute_exception_stack_depths: bool = True,

) -> _bytecode.ConcreteBytecode:

converter = _bytecode._ConvertBytecodeToConcrete(self)

return converter.to_concrete_bytecode(

compute_jumps_passes=compute_jumps_passes,

compute_exception_stack_depths=compute_exception_stack_depths,

)