Ensure stage not set warning is displayed correctly

lib/Capfile

@@ -1,2 +1,3 @@
 #!/usr/bin/env cap
+include Capistrano::DSL
 require 'capistrano/install'

lib/capistrano/application.rb

@@ -26,7 +26,7 @@ def top_level_tasks
       if tasks_without_stage_dependency.include?(@top_level_tasks.first)
         @top_level_tasks
       else
-        @top_level_tasks.unshift('deploy:ensure_stage')
+        @top_level_tasks.unshift(ensure_stage)
       end
     end
 
@@ -37,15 +37,6 @@ def capfile
       File.expand_path(File.join(File.dirname(__FILE__),'..','Capfile'))
     end
 
-    def tasks_without_stage_dependency
-      defined_stages = Dir['config/deploy/*.rb'].map { |f| File.basename(f, '.rb') }
-      defined_stages + default_tasks
-    end
-
-    def default_tasks
-      %w{install}
-    end
-
     def version
       ['--version', '-V',
        "Display the program version.",

lib/capistrano/defaults.rb

@@ -11,11 +11,3 @@
 
 set :pty, true
 
-namespace :deploy do
-  task :ensure_stage do
-    unless stage_set?
-      puts t(:stage_not_set)
-      exit 1
-    end
-  end
-end

lib/capistrano/dsl/task_enhancements.rb

@@ -32,5 +32,22 @@ def define_remote_file_task(task, target_roles)
       end
     end
 
+    def ensure_stage
+      Rake::Task.define_task(:ensure_stage) do
+        unless stage_set?
+          puts t(:stage_not_set)
+          exit 1
+        end
+      end
+    end
+
+    def tasks_without_stage_dependency
+      stages + default_tasks
+    end
+
+    def default_tasks
+      %w{install}
+    end
+
   end
 end

lib/capistrano/tasks/framework.rake

@@ -44,8 +44,6 @@ namespace :deploy do
   task :finished do
   end
 
-  before :starting, :ensure_stage
-
   desc 'Rollback to previous release.'
   task :rollback do
     %w{ starting started