set :stages, %w( vagrant staging production )
set :default_stage, "vagrant"
require 'capistrano/ext/multistage'

## === Replace this line ===
# env = ENV['RUBBER_ENV'] ||= ...

## === With the following code ===
## So that we can do `cap staging deploy` instead of RUBBER_ENV=staging cap deploy
env = ARGV.shift
# Default to vagrant if not provided
stages = %w( vagrant staging production )
default_stage = 'vagrant'
unless stages.include?(env)
  # using default stage, put that argument back
  ARGV.unshift(env)
  env = default_stage
end
ENV['RUBBER_ENV'] = ENV['RAILS_ENV'] = env