1# (C) Copyright 2020 by Rocky Bernstein
2#
3#  This program is free software; you can redistribute it and/or
4#  modify it under the terms of the GNU General Public License
5#  as published by the Free Software Foundation; either version 2
6#  of the License, or (at your option) any later version.
7#
8#  This program is distributed in the hope that it will be useful,
9#  but WITHOUT ANY WARRANTY; without even the implied warranty of
10#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11#  GNU General Public License for more details.
12#
13#  You should have received a copy of the GNU General Public License
14#  along with this program; if not, write to the Free Software
15#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
16
17import types
18from copy import deepcopy
19
20from xdis.codetype.code20 import Code2, Code2FieldTypes
21from xdis.version_info import PYTHON_VERSION
22
23# Below, in the Python 2.4 branch "bytes" is "str" since there may not be a "bytes" type.
24Code3FieldTypes = deepcopy(Code2FieldTypes)
25Code3FieldTypes.update({
26    "co_kwonlyargcount": int,
27})
28
29
30class Code3(Code2):
31    """Class for a Python3 code object used when a Python interpreter less than 3 is
32    working on Python3 bytecode. It also functions as an object that can be used
33    to build or write a Python3 code object, since we allow mutable structures.
34    When done mutating, call method freeze().
35
36    For convenience in generating code objects, fields like
37    `co_consts`, co_names which are (immutable) tuples in the end-result can be stored
38    instead as (mutable) lists. Likewise the line number table `co_lnotab`
39    can be stored as a simple list of offset, line_number tuples.
40    """
41
42    def __init__(
43        self,
44        co_argcount,
45        co_kwonlyargcount,
46        co_nlocals,
47        co_stacksize,
48        co_flags,
49        co_code,
50        co_consts,
51        co_names,
52        co_varnames,
53        co_filename,
54        co_name,
55        co_firstlineno,
56        co_lnotab,
57        co_freevars,
58        co_cellvars,
59    ):
60        super(Code3, self).__init__(
61            co_argcount,
62            co_nlocals,
63            co_stacksize,
64            co_flags,
65            co_code,
66            co_consts,
67            co_names,
68            co_varnames,
69            co_filename,
70            co_name,
71            co_firstlineno,
72            co_lnotab,
73            co_freevars,
74            co_cellvars,
75        )
76        self.co_kwonlyargcount = co_kwonlyargcount
77        self.fieldtypes = Code3FieldTypes
78
79        if type(self) == Code3:
80            self.check()
81        return
82
83    def encode_lineno_tab(self):
84        co_lnotab = b""
85
86        prev_line_number = self.co_firstlineno
87        prev_offset = 0
88        for offset, line_number in self.co_lnotab:
89            offset_diff = offset - prev_offset
90            line_diff = line_number - prev_line_number
91            prev_offset = offset
92            prev_line_number = line_number
93            while offset_diff >= 256:
94                co_lnotab += bytearray([255, 0])
95                offset_diff -= 255
96            while line_diff >= 256:
97                co_lnotab += bytearray([0, 255])
98                line_diff -= 255
99            co_lnotab += bytearray([offset_diff, line_diff])
100
101        self.co_lnotab = co_lnotab
102
103    def freeze(self):
104        for field in "co_consts co_names co_varnames co_freevars co_cellvars".split():
105            val = getattr(self, field)
106            if isinstance(val, list):
107                setattr(self, field, tuple(val))
108
109        # for field, typename in self.fieldtypes:
110        #     pass
111
112        if isinstance(self.co_lnotab, dict):
113            d = self.co_lnotab
114            self.co_lnotab = sorted(zip(d.keys(), d.values()), key=lambda tup: tup[0])
115        if isinstance(self.co_lnotab, list):
116            # We assume we have a list of tuples:
117            # (offset, linenumber) which we convert
118            # into the encoded format
119            self.encode_lineno_tab()
120
121        if isinstance(self.co_code, str) and PYTHON_VERSION >= 3.0:
122            self.co_code = self.co_code.encode()
123
124        if isinstance(self.co_lnotab, str):
125            self.co_lnotab = self.co_lnotab.encode()
126
127        return self
128
129    def to_native(self):
130        if not (3.0 <= PYTHON_VERSION <= 3.7):
131            raise TypeError(
132                "Python Interpreter needs to be in range 3.0..3.7; is %s"
133                % PYTHON_VERSION
134            )
135
136        code = deepcopy(self)
137        code.freeze()
138        try:
139            code.check()
140        except AssertionError as e:
141            raise TypeError(e)
142
143        return types.CodeType(
144            code.co_argcount,
145            code.co_kwonlyargcount,
146            code.co_nlocals,
147            code.co_stacksize,
148            code.co_flags,
149            code.co_code,
150            code.co_consts,
151            code.co_names,
152            code.co_varnames,
153            code.co_filename,
154            code.co_name,
155            code.co_firstlineno,
156            code.co_lnotab,
157            code.co_freevars,
158            code.co_cellvars,
159        )
160