Fix #partition_union

soutaro · Jul 6, 2023 · 7d85557 · 7d85557
1 parent b1d0fcc
commit 7d85557
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/lib/steep/ast/types/factory.rb b/lib/steep/ast/types/factory.rb
@@ -358,16 +358,21 @@ def partition_union(type)
               partition_union(unfold)
             end
           when AST::Types::Union
-            falsy_types, truthy_types = type.types.partition do |type|
-              (type.is_a?(AST::Types::Literal) && type.value == false) ||
-                type.is_a?(AST::Types::Nil)
+            truthy_types = [] #: Array[AST::Types::t]
+            falsy_types = [] #: Array[AST::Types::t]
+
+            type.types.each do |type|
+              truthy, falsy = partition_union(type)
+
+              truthy_types << truthy if truthy
+              falsy_types << falsy if falsy
             end
 
             [
               truthy_types.empty? ? nil : AST::Types::Union.build(types: truthy_types),
               falsy_types.empty? ? nil : AST::Types::Union.build(types: falsy_types)
             ]
-          when AST::Types::Any, AST::Types::Boolean
+          when AST::Types::Any, AST::Types::Boolean, AST::Types::Top
             [type, type]
           when AST::Types::Nil
             [nil, type]

diff --git a/test/type_factory_test.rb b/test/type_factory_test.rb
@@ -395,4 +395,31 @@ class Bar
       end
     end
   end
+
+  def test_partition_union__top
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("top"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("top")), truthy
+        assert_equal factory.type(parse_type("top")), falsy
+      end
+    end
+  end
+
+  def test_partition_union__boolish
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("::boolish"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("top")), truthy
+        assert_equal factory.type(parse_type("top")), falsy
+      end
+    end
+  end
+
+  def test_partition_union__bool_union
+    with_factory() do |factory|
+      factory.partition_union(factory.type(parse_type("bool | ::Symbol"))).tap do |truthy, falsy|
+        assert_equal factory.type(parse_type("bool | ::Symbol")), truthy
+        assert_equal factory.type(parse_type("bool")), falsy
+      end
+    end
+  end
 end