1 files changed, 6 insertions, 13 deletions
diff --git a/lib/puppet/parser/compiler.rb b/lib/puppet/parser/compiler.rb
index 98cad952c..7504b276b 100644
--- a/lib/puppet/parser/compiler.rb
+++ b/lib/puppet/parser/compiler.rb
@@ -56,13 +56,10 @@ class Puppet::Parser::Compiler
     # Note that this will fail if the resource is not unique.
     @catalog.add_resource(resource)
 
-    set_container_resource(scope, resource)
-  end
 
-  # Add our container edge.  If we're a class, then we get treated specially - we can
-  # control the stage that the class is applied in.  Otherwise, we just
-  # get added to our parent container.
-  def set_container_resource(scope, resource)
+    # Add our container edge.  If we're a class, then we get treated specially - we can
+    # control the stage that the class is applied in.  Otherwise, we just
+    # get added to our parent container.
     return if resource.type.to_s.downcase == "stage"
 
     if resource.type.to_s.downcase != "class"
@@ -70,15 +67,14 @@ class Puppet::Parser::Compiler
       return @catalog.add_edge(scope.resource, resource)
     end
 
-    unless stage = @catalog.resource(:stage, resource[:stage] || :main)
+    unless stage = @catalog.resource(:stage, resource[:stage] || (scope && scope.resource && scope.resource[:stage]) || :main)
       raise ArgumentError, "Could not find stage #{resource[:stage] || :main} specified by #{resource}"
     end
 
+    resource[:stage] ||= stage.title
     @catalog.add_edge(stage, resource)
   end
 
-  private :set_container_resource
-
   # Do we use nodes found in the code, vs. the external node sources?
   def ast_nodes?
     known_resource_types.nodes?
@@ -289,10 +285,7 @@ class Puppet::Parser::Compiler
     @main_resource = Puppet::Parser::Resource.new("class", :main, :scope => @topscope, :source => @main)
     @topscope.resource = @main_resource
 
-    @resources << @main_resource
-    @catalog.add_resource(@main_resource)
-
-    set_container_resource(@topscope, @main_resource)
+    add_resource(@topscope, @main_resource)
 
     @main_resource.evaluate
   end