From 58017d02987e9fdbae23750c3ab1079f47b3e5c5 Mon Sep 17 00:00:00 2001
From: Dawid Marcin Grzesiak <dmg@dmg.pl>
Date: Mon, 29 Mar 2010 04:06:10 +0200
Subject: [PATCH] possibility to create custom discriminator types

---
 lib/dm-core/model.rb                    | 2 +-
 lib/dm-core/property_set.rb             | 2 +-
 spec/public/types/discriminator_spec.rb | 8 +++++++-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/lib/dm-core/model.rb b/lib/dm-core/model.rb
index 30ccce9e..eb0283f7 100644
--- a/lib/dm-core/model.rb
+++ b/lib/dm-core/model.rb
@@ -548,7 +548,7 @@ def load(records, query)
             record = record.dup
             field_map.each { |property, field| record[property] = record.delete(field) if record.key?(field) }
 
-            model     = discriminator && record[discriminator] || self
+            model     = discriminator && (klass = record[discriminator]) && discriminator.typecast(klass) || self
             model_key = model.key(repository_name)
 
             resource = if model_key.valid?(key_values = record.values_at(*model_key))
diff --git a/lib/dm-core/property_set.rb b/lib/dm-core/property_set.rb
index 1300a4b8..8bb37fec 100644
--- a/lib/dm-core/property_set.rb
+++ b/lib/dm-core/property_set.rb
@@ -68,7 +68,7 @@ def key
 
     # @api semipublic
     def discriminator
-      @discriminator ||= detect { |property| property.type == Types::Discriminator }
+      @discriminator ||= detect { |property| property.type.ancestors.include?(Types::Discriminator) }
     end
 
     # @api semipublic
diff --git a/spec/public/types/discriminator_spec.rb b/spec/public/types/discriminator_spec.rb
index cefce8c8..786c8442 100644
--- a/spec/public/types/discriminator_spec.rb
+++ b/spec/public/types/discriminator_spec.rb
@@ -3,12 +3,18 @@
 describe DataMapper::Types::Discriminator do
   before :all do
     module ::Blog
+      class NewDiscriminator < DataMapper::Types::Discriminator
+        primitive Class
+        default   lambda { |resource, property| resource.model }
+        required  true
+      end
+
       class Article
         include DataMapper::Resource
 
         property :id,    Serial
         property :title, String, :required => true
-        property :type,  Discriminator
+        property :type,  NewDiscriminator
       end
 
       class Announcement < Article; end