Custom 404

lib/sinatra.rb

@@ -27,7 +27,7 @@ end
 module Sinatra
   extend self
 
-  Result = Struct.new(:block, :params)
+  Result = Struct.new(:block, :params, :status) unless defined?(Result)
   
   def application
     @app ||= Application.new
@@ -83,12 +83,24 @@ module Sinatra
     def invoke(env)
       return unless pattern =~ env['PATH_INFO'].squeeze('/')
       params = param_keys.zip($~.captures.map(&:from_param)).to_hash
-      Result.new(block, params)
+      Result.new(block, params, 200)
     end
     
   end
   
-  class EventContext
+  class Error
+    
+    attr_reader :code, :block
+    
+    def initialize(code, &b)
+      @code, @block = code, b
+    end
+    
+    def invoke(env)
+      Result.new(block, {}, 404)
+    end
+    
+  end
   
   module ResponseHelpers
 
@@ -97,7 +109,6 @@ module Sinatra
     end
 
   end
-    include ResponseHelpers
   
   module RenderingHelpers
     
@@ -154,6 +165,10 @@ module Sinatra
       end
     
   end
+
+  class EventContext
+    
+    include ResponseHelpers
     include RenderingHelpers
     
     attr_accessor :request, :response
@@ -229,8 +244,19 @@ module Sinatra
       layouts[name] = b
     end
     
+    def define_error(code, &b)
+      events[:errors][code] = Error.new(code, &b)
+    end
+    
     def lookup(env)
-      events[env['REQUEST_METHOD'].downcase.to_sym].eject(&[:invoke, env])
+      e = events[env['REQUEST_METHOD'].downcase.to_sym].eject(&[:invoke, env])
+      e ||= (events[:errors][404] || basic_not_found).invoke(env)
+    end
+    
+    def basic_not_found
+      Error.new(404) do
+        '<h1>Not Found</h1>'
+      end
     end
 
     def options
@@ -238,12 +264,13 @@ module Sinatra
     end
     
     def call(env)
-      return [404, {}, 'Not Found'] unless result = lookup(env)
+      result = lookup(env)
       context = EventContext.new(
         Rack::Request.new(env), 
         Rack::Response.new,
         result.params
       )
+      context.status(result.status)
       returned = catch(:halt) do
         [:complete, context.instance_eval(&result.block)]
       end
@@ -276,6 +303,10 @@ def helpers(&b)
   Sinatra::EventContext.class_eval(&b)
 end
 
+def error(code, &b)
+  Sinatra.application.define_error(code, &b)
+end
+
 def layout(name = :layout, &b)
   Sinatra.application.define_layout(name, &b)
 end

test/application_test.rb

@@ -45,7 +45,7 @@ context "An app returns" do
     request = Rack::MockRequest.new(@app)
     result = request.get('/')
     result.should.be.not_found
-    result.body.should.equal 'Not Found'
+    result.body.should.equal '<h1>Not Found</h1>'
   end
   
   specify "200 if success" do