Add two asserts to the ast snippets to check attributes of ast nodes.

Author: windelbouwman

tests/snippets/ast_snippet.py

@@ -7,3 +7,10 @@ def foo():
 """
 n = ast.parse(source)
 print(n)
+print(n.body)
+print(n.body[0].name)
+assert n.body[0].name == 'foo'
+print(n.body[0].body)
+print(n.body[0].body[0])
+print(n.body[0].body[0].value.func.id)
+assert n.body[0].body[0].value.func.id == 'print'

vm/src/pyobject.rs

@@ -213,6 +213,15 @@ impl PyContext {
         self.object.clone()
     }
 
+    pub fn new_object(&self) -> PyObjectRef {
+        PyObject::new(
+            PyObjectKind::Instance {
+                dict: self.new_dict(),
+            },
+            self.object(),
+        )
+    }
+
     pub fn new_int(&self, i: i32) -> PyObjectRef {
         PyObject::new(PyObjectKind::Integer { value: i }, self.int_type())
     }

vm/src/stdlib/ast.rs

@@ -8,10 +8,9 @@
 extern crate rustpython_parser;
 
 use self::rustpython_parser::{ast, parser};
-use super::super::obj::{objdict, objfloat, objint, objlist, objstr, objtuple, objtype};
-use super::super::objbool;
+use super::super::obj::{objstr, objtype};
 use super::super::pyobject::{
-    DictProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::VirtualMachine;
 
@@ -20,31 +19,92 @@ fn program_to_ast(ctx: &PyContext, program: &ast::Program) -> PyObjectRef {
     for statement in &program.statements {
         body.push(statement_to_ast(ctx, statement));
     }
-    // TODO: create Module node and set attributes:
-    let ast_node = ctx.new_list(body);
+    // TODO: create Module node:
+    // let ast_node = ctx.new_instance(this.Module);
+    let ast_node = ctx.new_object();
+    let py_body = ctx.new_list(body);
+    ast_node.set_attr("body", py_body);
     ast_node
 }
 
 fn statement_to_ast(ctx: &PyContext, statement: &ast::LocatedStatement) -> PyObjectRef {
-    match &statement.node {
-        ast::Statement::FunctionDef { name, args, body } => {
-            // TODO: create ast.FunctionDef object and set attributes
-            // let node = ctx.new_object();
+    let node = match &statement.node {
+        ast::Statement::FunctionDef {
+            name,
+            args: _,
+            body,
+        } => {
+            // TODO: create ast.FunctionDef object:
+            let node = ctx.new_object();
+
+            // Set name:
+            node.set_attr("name", ctx.new_str(name.to_string()));
+
+            // Set body:
             let mut py_body = vec![];
-            py_body.push(ctx.new_str(format!("{:?}", name)));
-            // node.set_attr("name", new_str(name));
             for statement in body {
                 py_body.push(statement_to_ast(ctx, statement));
             }
 
-            ctx.new_list(py_body)
-        },
-        _ => {
-            ctx.new_str(format!("{:?}", statement))
+            node.set_attr("body", ctx.new_list(py_body));
+            node
         }
-    }
+        ast::Statement::Expression { expression } => {
+            let value = expression_to_ast(ctx, expression);
+            // TODO: create proper class:
+            let node = ctx.new_object();
+            node.set_attr("value", value);
+            node
+        }
+        x => {
+            unimplemented!("{:?}", x);
+        }
+    };
+
+    // set lineno on node:
+    let lineno = ctx.new_int(statement.location.get_row() as i32);
+    node.set_attr("lineno", lineno);
+
+    node
+}
+
+fn expression_to_ast(ctx: &PyContext, expression: &ast::Expression) -> PyObjectRef {
+    let node = match &expression {
+        ast::Expression::Call { function, args } => {
+            // TODO: create ast.Call instance
+            let node = ctx.new_object();
+
+            let py_func_ast = expression_to_ast(ctx, function);
+            node.set_attr("func", py_func_ast);
+
+            let mut py_args = vec![];
+            for arg in args {
+                py_args.push(expression_to_ast(ctx, &arg.1));
+            }
+            let py_args = ctx.new_list(py_args);
+            node.set_attr("args", py_args);
+
+            node
+        }
+        ast::Expression::Identifier { name } => {
+            // TODO: create ast.Identifier instance
+            let node = ctx.new_object();
+            let py_name = ctx.new_str(name.clone());
+            node.set_attr("id", py_name);
+            node
+        }
+        ast::Expression::String { value } => {
+            let node = ctx.new_object();
+            node.set_attr("s", ctx.new_str(value.clone()));
+            node
+        }
+        n => {
+            unimplemented!("{:?}", n);
+        }
+    };
 
     // TODO: set lineno on object
+    node
 }
 
 fn ast_parse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -66,5 +126,18 @@ fn ast_parse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     let ast_mod = ctx.new_module(&"ast".to_string(), ctx.new_scope(None));
     ast_mod.set_item("parse", ctx.new_rustfunc(ast_parse));
+    ast_mod.set_item(
+        "Module",
+        ctx.new_class(&"_ast.Module".to_string(), ctx.object()),
+    );
+    // TODO: maybe we can use some clever macro to generate this?
+    ast_mod.set_item(
+        "FunctionDef",
+        ctx.new_class(&"_ast.FunctionDef".to_string(), ctx.object()),
+    );
+    ast_mod.set_item(
+        "Call",
+        ctx.new_class(&"_ast.Call".to_string(), ctx.object()),
+    );
     ast_mod
 }