Fix for #4542 -- included classes weren't assigned proper stages

This commit unifies the code paths on which classes are added, alters the default
stage to respect the stage of the parent if any, and assures that the resource is
notified if its stage is assigned (turning an implicit stage into an explicit one).

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