From 96b1857d0bf0015d05a632ef4e35a38bd91ec717 Mon Sep 17 00:00:00 2001
From: seenmyfate <seenmyfate@gmail.com>
Date: Fri, 6 Sep 2013 13:12:34 +0100
Subject: [PATCH] Ensure `stage not set` warning is displayed correctly

---
 lib/Capfile                             |  1 +
 lib/capistrano/application.rb           | 11 +----------
 lib/capistrano/defaults.rb              |  8 --------
 lib/capistrano/dsl/task_enhancements.rb | 17 +++++++++++++++++
 lib/capistrano/tasks/framework.rake     |  2 --
 5 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/lib/Capfile b/lib/Capfile
index 6a8e2e93..8211354e 100644
--- a/lib/Capfile
+++ b/lib/Capfile
@@ -1,2 +1,3 @@
 #!/usr/bin/env cap
+include Capistrano::DSL
 require 'capistrano/install'
diff --git a/lib/capistrano/application.rb b/lib/capistrano/application.rb
index ec2adca8..4f2a67bb 100644
--- a/lib/capistrano/application.rb
+++ b/lib/capistrano/application.rb
@@ -26,7 +26,7 @@ module Capistrano
       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 @@ module Capistrano
       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.",
diff --git a/lib/capistrano/defaults.rb b/lib/capistrano/defaults.rb
index 36d92c03..0c105d0e 100644
--- a/lib/capistrano/defaults.rb
+++ b/lib/capistrano/defaults.rb
@@ -11,11 +11,3 @@ set :log_level, :debug
 
 set :pty, true
 
-namespace :deploy do
-  task :ensure_stage do
-    unless stage_set?
-      puts t(:stage_not_set)
-      exit 1
-    end
-  end
-end
diff --git a/lib/capistrano/dsl/task_enhancements.rb b/lib/capistrano/dsl/task_enhancements.rb
index a3536adb..29090652 100644
--- a/lib/capistrano/dsl/task_enhancements.rb
+++ b/lib/capistrano/dsl/task_enhancements.rb
@@ -32,5 +32,22 @@ module Capistrano
       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
diff --git a/lib/capistrano/tasks/framework.rake b/lib/capistrano/tasks/framework.rake
index 97770b53..6ff02eef 100644
--- a/lib/capistrano/tasks/framework.rake
+++ b/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