Enum from_value when Flags (#3518)

crystal-lang · Nov 20, 2016 · 584060e · 584060e
1 parent d3bc110
commit 584060e
Showing 2 changed files with 39 additions and 12 deletions.
diff --git a/spec/std/enum_spec.cr b/spec/std/enum_spec.cr
@@ -83,18 +83,37 @@ describe Enum do
     end
   end
 
-  it "does from_value?" do
-    SpecEnum.from_value?(0).should eq(SpecEnum::One)
-    SpecEnum.from_value?(1).should eq(SpecEnum::Two)
-    SpecEnum.from_value?(2).should eq(SpecEnum::Three)
-    SpecEnum.from_value?(3).should be_nil
+  describe "from_value?" do
+    it "for simple enum" do
+      SpecEnum.from_value?(0).should eq(SpecEnum::One)
+      SpecEnum.from_value?(1).should eq(SpecEnum::Two)
+      SpecEnum.from_value?(2).should eq(SpecEnum::Three)
+      SpecEnum.from_value?(3).should be_nil
+    end
+
+    it "for flags enum" do
+      SpecEnumFlags.from_value?(0).should be_nil
+      SpecEnumFlags.from_value?(1).should eq(SpecEnumFlags::One)
+      SpecEnumFlags.from_value?(2).should eq(SpecEnumFlags::Two)
+      SpecEnumFlags.from_value?(3).should eq(SpecEnumFlags::One | SpecEnumFlags::Two)
+      SpecEnumFlags.from_value?(8).should be_nil
+    end
   end
 
-  it "does from_value" do
-    SpecEnum.from_value(0).should eq(SpecEnum::One)
-    SpecEnum.from_value(1).should eq(SpecEnum::Two)
-    SpecEnum.from_value(2).should eq(SpecEnum::Three)
-    expect_raises { SpecEnum.from_value(3) }
+  describe "from_value" do
+    it "for simple enum" do
+      SpecEnum.from_value(0).should eq(SpecEnum::One)
+      SpecEnum.from_value(1).should eq(SpecEnum::Two)
+      SpecEnum.from_value(2).should eq(SpecEnum::Three)
+      expect_raises { SpecEnum.from_value(3) }
+    end
+
+    it "for flags enum" do
+      expect_raises { SpecEnumFlags.from_value(0) }
+      SpecEnumFlags.from_value(1).should eq(SpecEnumFlags::One)
+      SpecEnumFlags.from_value(2).should eq(SpecEnumFlags::Two)
+      SpecEnumFlags.from_value(3).should eq(SpecEnumFlags::One | SpecEnumFlags::Two)
+    end
   end
 
   it "has hash" do

diff --git a/src/enum.cr b/src/enum.cr
@@ -310,8 +310,16 @@ struct Enum
   # Color.from_value?(3) # => nil
   # ```
   def self.from_value?(value) : self | Nil
-    {% for member in @type.constants %}
-      return {{@type}}::{{member}} if {{@type}}::{{member}}.value == value
+    {% if @type.has_attribute?("Flags") %}
+      mask = {% for member, i in @type.constants %}\
+        {% if i != 0 %} | {% end %}\
+        {{@type}}::{{member}}.value{% end %}
+      return if (mask & value != value) || (value == 0 && values.none? { |val| val.to_i == 0 })
+      return new(value)
+    {% else %}
+      {% for member in @type.constants %}
+        return {{@type}}::{{member}} if {{@type}}::{{member}}.value == value
+      {% end %}
     {% end %}
     nil
   end