4 files changed, 21 insertions, 9 deletions
diff --git a/lib/puppet/parser/ast/definition.rb b/lib/puppet/parser/ast/definition.rb
index 287845ade..c43422f82 100644
--- a/lib/puppet/parser/ast/definition.rb
+++ b/lib/puppet/parser/ast/definition.rb
@@ -3,12 +3,15 @@ require 'puppet/parser/ast/top_level_construct'
 class Puppet::Parser::AST::Definition < Puppet::Parser::AST::TopLevelConstruct
   attr_accessor :context
 
-  def initialize(name, context = {})
+  def initialize(name, context = {}, &ruby_code)
     @name = name
     @context = context
+    @ruby_code = ruby_code
   end
 
   def instantiate(modname)
-    return [Puppet::Resource::Type.new(:definition, @name, @context.merge(:module_name => modname))]
+    new_definition = Puppet::Resource::Type.new(:definition, @name, @context.merge(:module_name => modname))
+    new_definition.ruby_code = @ruby_code if @ruby_code
+    [new_definition]
   end
 end
diff --git a/lib/puppet/parser/ast/hostclass.rb b/lib/puppet/parser/ast/hostclass.rb
index d539e4deb..cab5e4a24 100644
--- a/lib/puppet/parser/ast/hostclass.rb
+++ b/lib/puppet/parser/ast/hostclass.rb
@@ -3,13 +3,16 @@ require 'puppet/parser/ast/top_level_construct'
 class Puppet::Parser::AST::Hostclass < Puppet::Parser::AST::TopLevelConstruct
   attr_accessor :name, :context
 
-  def initialize(name, context = {})
+  def initialize(name, context = {}, &ruby_code)
     @context = context
     @name = name
+    @ruby_code = ruby_code
   end
 
   def instantiate(modname)
-    all_types = [Puppet::Resource::Type.new(:hostclass, @name, @context.merge(:module_name => modname))]
+    new_class = Puppet::Resource::Type.new(:hostclass, @name, @context.merge(:module_name => modname))
+    new_class.ruby_code = @ruby_code if @ruby_code
+    all_types = [new_class]
     if code
       code.each do |nested_ast_node|
         if nested_ast_node.respond_to? :instantiate
diff --git a/lib/puppet/parser/ast/node.rb b/lib/puppet/parser/ast/node.rb
index 4951a6365..b69a5c4e0 100644
--- a/lib/puppet/parser/ast/node.rb
+++ b/lib/puppet/parser/ast/node.rb
@@ -3,15 +3,18 @@ require 'puppet/parser/ast/top_level_construct'
 class Puppet::Parser::AST::Node < Puppet::Parser::AST::TopLevelConstruct
   attr_accessor :names, :context
 
-  def initialize(names, context = {})
+  def initialize(names, context = {}, &ruby_code)
     raise ArgumentError, "names should be an array" unless names.is_a? Array
     @names = names
     @context = context
+    @ruby_code = ruby_code
   end
 
   def instantiate(modname)
     @names.collect do |name|
-      Puppet::Resource::Type.new(:node, name, @context.merge(:module_name => modname))
+      new_node = Puppet::Resource::Type.new(:node, name, @context.merge(:module_name => modname))
+      new_node.ruby_code = @ruby_code if @ruby_code
+      new_node
     end
   end
 end
diff --git a/lib/puppet/parser/parser_support.rb b/lib/puppet/parser/parser_support.rb
index 288abb582..16cd0ee0f 100644
--- a/lib/puppet/parser/parser_support.rb
+++ b/lib/puppet/parser/parser_support.rb
@@ -146,8 +146,7 @@ class Puppet::Parser::Parser
   # how should I do error handling here?
   def parse(string = nil)
     if self.file =~ /\.rb$/
-      parse_ruby_file
-      main = nil
+      main = parse_ruby_file
     else
       self.string = string if string
       begin
@@ -189,7 +188,11 @@ class Puppet::Parser::Parser
   def parse_ruby_file
     # Execute the contents of the file inside its own "main" object so
     # that it can call methods in the resource type API.
-    Puppet::DSL::ResourceTypeAPI.new.instance_eval(File.read(self.file))
+    main_object = Puppet::DSL::ResourceTypeAPI.new
+    main_object.instance_eval(File.read(self.file))
+
+    # Then extract any types that were created.
+    Puppet::Parser::AST::ASTArray.new :children => main_object.instance_eval { @__created_ast_objects__ }
   end
 
   def string=(string)