Skip to content

Commit d1e5066

Browse files
cosminbascacosminbasca
committed
Code object fix missing attributes for the generator case.
1 parent ed72c0c commit d1e5066

File tree

4 files changed

+53
-37
lines changed

4 files changed

+53
-37
lines changed

‎graalpython/com.oracle.graal.python.test/src/tests/test_code.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ def inner_func():
5555

5656
try:
5757
loc_1 &= kwarg_other
58+
yield loc_1
5859
except TypeError:
5960
pass
6061
else:
@@ -81,12 +82,14 @@ def test_code_attributes():
8182
assert code.co_kwonlyargcount == 0
8283
assert code.co_nlocals == 6
8384
assert code.co_stacksize >= code.co_nlocals
84-
assert code.co_flags == 0
85+
assert code.co_flags & (1 << 5)
86+
assert not code.co_flags & (1 << 2)
87+
assert not code.co_flags & (1 << 3)
8588
# assert code.co_code
8689
# assert code.co_consts
8790
# assert set(code.co_names) == {'set', 'TypeError', 'print'}
8891
assert set(code.co_varnames) == {'arg_l', 'kwarg_case', 'kwarg_other', 'loc_1', 'loc_3', 'inner_func'}
89-
assert code.co_filename == "graalpython/com.oracle.graal.python.test/src/tests/test_code.py"
92+
assert code.co_filename.endswith("test_code.py")
9093
assert code.co_name == "my_func"
9194
assert code.co_firstlineno == 48
9295
# assert code.co_lnotab == b'\x00\x01\x0c\x01\x0c\x01\x06\x02\x15\x03\x03\x01\x0e\x01\r\x01\x05\x02'

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/code/PCode.java

Lines changed: 39 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
4949
import com.oracle.graal.python.builtins.objects.type.PythonClass;
5050
import com.oracle.graal.python.nodes.ModuleRootNode;
51+
import com.oracle.graal.python.nodes.PNode;
5152
import com.oracle.graal.python.nodes.argument.ReadIndexedArgumentNode;
5253
import com.oracle.graal.python.nodes.argument.ReadKeywordNode;
5354
import com.oracle.graal.python.nodes.argument.ReadVarArgsNode;
@@ -101,7 +102,6 @@ public class PCode extends PythonBuiltinObject {
101102
// tuple of names of cell variables (referenced by containing scopes)
102103
private Object[] cellvars;
103104

104-
@TruffleBoundary
105105
public PCode(PythonClass cls, RootNode rootNode, PythonCore core) {
106106
super(cls);
107107
this.rootNode = rootNode;
@@ -135,11 +135,12 @@ public PCode(PythonClass cls, int argcount, int kwonlyargcount,
135135
this.cellvars = cellvars;
136136
}
137137

138+
@TruffleBoundary
138139
private static Set<String> asSet(String[] values) {
139140
return (values != null) ? new HashSet<>(Arrays.asList(values)) : new HashSet<>();
140141
}
141142

142-
private static String[] getFreeVars(RootNode rootNode) {
143+
private static String[] extractFreeVars(RootNode rootNode) {
143144
if (rootNode instanceof FunctionRootNode) {
144145
return ((FunctionRootNode) rootNode).getFreeVars();
145146
} else if (rootNode instanceof GeneratorFunctionRootNode) {
@@ -149,7 +150,7 @@ private static String[] getFreeVars(RootNode rootNode) {
149150
}
150151
}
151152

152-
private static String[] getCellVars(RootNode rootNode) {
153+
private static String[] extractCellVars(RootNode rootNode) {
153154
if (rootNode instanceof FunctionRootNode) {
154155
return ((FunctionRootNode) rootNode).getCellVars();
155156
} else if (rootNode instanceof GeneratorFunctionRootNode) {
@@ -160,18 +161,20 @@ private static String[] getCellVars(RootNode rootNode) {
160161
}
161162

162163
private static String extractFileName(RootNode rootNode) {
163-
SourceSection src = rootNode.getSourceSection();
164+
RootNode funcRootNode = (rootNode instanceof GeneratorFunctionRootNode) ? ((GeneratorFunctionRootNode) rootNode).getFunctionRootNode() : rootNode;
165+
SourceSection src = funcRootNode.getSourceSection();
164166
if (src != null) {
165167
return src.getSource().getName();
166-
} else if (rootNode instanceof ModuleRootNode) {
167-
return rootNode.getName();
168+
} else if (funcRootNode instanceof ModuleRootNode) {
169+
return funcRootNode.getName();
168170
} else {
169171
return null;
170172
}
171173
}
172174

173175
private static int extractFirstLineno(RootNode rootNode) {
174-
SourceSection sourceSection = rootNode.getSourceSection();
176+
RootNode funcRootNode = (rootNode instanceof GeneratorFunctionRootNode) ? ((GeneratorFunctionRootNode) rootNode).getFunctionRootNode() : rootNode;
177+
SourceSection sourceSection = funcRootNode.getSourceSection();
175178
return (sourceSection != null) ? sourceSection.getStartLine() : 1;
176179
}
177180

@@ -187,6 +190,7 @@ private static String extractName(RootNode rootNode) {
187190
return name;
188191
}
189192

193+
@TruffleBoundary
190194
private static Set<String> getKeywordArgumentNames(List<ReadKeywordNode> readKeywordNodes) {
191195
Set<String> kwArgNames = new HashSet<>();
192196
for (ReadKeywordNode node : readKeywordNodes) {
@@ -195,9 +199,10 @@ private static Set<String> getKeywordArgumentNames(List<ReadKeywordNode> readKey
195199
return kwArgNames;
196200
}
197201

198-
private static Set<String> getArgumentNames(List<ReadIndexedArgumentNode> readIndexedArgumentNodes) {
202+
@TruffleBoundary
203+
private static Set<String> extractArgumentNames(List<? extends PNode> readIndexedArgumentNodes) {
199204
Set<String> argNames = new HashSet<>();
200-
for (ReadIndexedArgumentNode node : readIndexedArgumentNodes) {
205+
for (PNode node : readIndexedArgumentNodes) {
201206
Node parent = node.getParent();
202207
if (parent instanceof WriteIdentifierNode) {
203208
Object identifier = ((WriteIdentifierNode) parent).getIdentifier();
@@ -213,25 +218,42 @@ private static int extractStackSize(RootNode rootNode) {
213218
return rootNode.getFrameDescriptor().getSize();
214219
}
215220

221+
@TruffleBoundary
216222
private void extractArgStats() {
217-
this.freevars = getFreeVars(rootNode);
218-
this.cellvars = getCellVars(rootNode);
219-
Set<String> freeVarsSet = asSet((String[])freevars);
220-
Set<String> cellVarsSet = asSet((String[])cellvars);
223+
// 0x20 - generator
224+
this.flags = 0;
225+
RootNode funcRootNode = rootNode;
226+
if (funcRootNode instanceof GeneratorFunctionRootNode) {
227+
flags |= (1 << 5);
228+
funcRootNode = ((GeneratorFunctionRootNode) funcRootNode).getFunctionRootNode();
229+
}
221230

222-
List<ReadKeywordNode> readKeywordNodes = NodeUtil.findAllNodeInstances(rootNode, ReadKeywordNode.class);
223-
List<ReadIndexedArgumentNode> readIndexedArgumentNodes = NodeUtil.findAllNodeInstances(rootNode, ReadIndexedArgumentNode.class);
231+
// 0x04 - *arguments
232+
if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarArgsNode.class).size() == 1) {
233+
flags |= (1 << 2);
234+
}
235+
// 0x08 - **keywords
236+
if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarKeywordsNode.class).size() == 1) {
237+
flags |= (1 << 3);
238+
}
239+
240+
this.freevars = extractFreeVars(rootNode);
241+
this.cellvars = extractCellVars(rootNode);
242+
Set<String> freeVarsSet = asSet((String[]) freevars);
243+
Set<String> cellVarsSet = asSet((String[]) cellvars);
244+
245+
List<ReadKeywordNode> readKeywordNodes = NodeUtil.findAllNodeInstances(funcRootNode, ReadKeywordNode.class);
246+
List<ReadIndexedArgumentNode> readIndexedArgumentNodes = NodeUtil.findAllNodeInstances(funcRootNode, ReadIndexedArgumentNode.class);
224247

225248
Set<String> kwNames = getKeywordArgumentNames(readKeywordNodes);
226-
Set<String> argNames = getArgumentNames(readIndexedArgumentNodes);
249+
Set<String> argNames = extractArgumentNames(readIndexedArgumentNodes);
227250

228251
Set<String> allArgNames = new HashSet<>();
229252
allArgNames.addAll(kwNames);
230253
allArgNames.addAll(argNames);
231254

232255
this.argcount = readIndexedArgumentNodes.size();
233256
this.kwonlyargcount = 0;
234-
this.flags = 0;
235257

236258
for (ReadKeywordNode kwNode : readKeywordNodes) {
237259
if (!kwNode.canBePositional()) {
@@ -254,22 +276,6 @@ private void extractArgStats() {
254276
}
255277
}
256278

257-
// 0x20 - generator
258-
RootNode funcRootNode = rootNode;
259-
if (funcRootNode instanceof GeneratorFunctionRootNode) {
260-
flags |= (1 << 5);
261-
funcRootNode = ((GeneratorFunctionRootNode) funcRootNode).getFunctionRootNode();
262-
}
263-
264-
// 0x04 - *arguments
265-
if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarArgsNode.class).size() == 1) {
266-
flags |= (1 << 2);
267-
}
268-
// 0x08 - **keywords
269-
if (NodeUtil.findAllNodeInstances(funcRootNode, ReadVarKeywordsNode.class).size() == 1) {
270-
flags |= (1 << 3);
271-
}
272-
273279
this.varnames = varnamesSet.toArray();
274280
this.nlocals = varnamesSet.size();
275281
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/function/PGeneratorFunction.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public static PGeneratorFunction create(PythonClass clazz, PythonCore core, Stri
4141
FrameDescriptor frameDescriptor, PythonObject globals, PCell[] closure, ExecutionCellSlots executionCellSlots,
4242
int numOfActiveFlags, int numOfGeneratorBlockNode, int numOfGeneratorForNode) {
4343

44-
GeneratorFunctionRootNode generatorFunctionRootNode = new GeneratorFunctionRootNode(core.getLanguage(), callTarget,
44+
GeneratorFunctionRootNode generatorFunctionRootNode = new GeneratorFunctionRootNode(core.getLanguage(), callTarget, name,
4545
frameDescriptor, closure, executionCellSlots, numOfActiveFlags, numOfGeneratorBlockNode, numOfGeneratorForNode);
4646

4747
return new PGeneratorFunction(clazz, name, enclosingClassName, arity, Truffle.getRuntime().createCallTarget(generatorFunctionRootNode), frameDescriptor, globals, closure);

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/generator/GeneratorFunctionRootNode.java

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,15 @@ public class GeneratorFunctionRootNode extends PClosureFunctionRootNode {
5858
private final int numOfGeneratorForNode;
5959
private final PCell[] closure;
6060
private final ExecutionCellSlots cellSlots;
61+
private final String name;
6162

6263
@Child private PythonObjectFactory factory = PythonObjectFactory.create();
6364

64-
public GeneratorFunctionRootNode(PythonLanguage language, RootCallTarget callTarget, FrameDescriptor frameDescriptor, PCell[] closure, ExecutionCellSlots executionCellSlots,
65+
public GeneratorFunctionRootNode(PythonLanguage language, RootCallTarget callTarget, String name, FrameDescriptor frameDescriptor, PCell[] closure, ExecutionCellSlots executionCellSlots,
6566
int numOfActiveFlags, int numOfGeneratorBlockNode, int numOfGeneratorForNode) {
6667
super(language, frameDescriptor, executionCellSlots);
6768
this.callTarget = callTarget;
69+
this.name = name;
6870
this.frameDescriptor = frameDescriptor;
6971
this.closure = closure;
7072
this.cellSlots = executionCellSlots;
@@ -81,4 +83,9 @@ public Object execute(VirtualFrame frame) {
8183
public RootNode getFunctionRootNode() {
8284
return callTarget.getRootNode();
8385
}
86+
87+
@Override
88+
public String getName() {
89+
return name;
90+
}
8491
}

0 commit comments

Comments
 (0)