diff --git a/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala b/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala
index 42e786b..b8b0d81 100644
--- a/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala
+++ b/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala
@@ -162,7 +162,7 @@ class Narrow(pipelineContext: PipelineContext) {
     t match {
       case DictMap(kt, _) => kt
       case ShapeMap(Nil)  => NoneType
-      case ShapeMap(_)    => AtomType
+      case ShapeMap(kvs)  => UnionType(Set(kvs.map(kv => AtomLitType(kv.key)): _*))
       case UnionType(ts)  => subtype.join(ts.map(getKeyType))
       case NoneType       => NoneType
       case _              => throw new IllegalStateException()
diff --git a/eqwalizer/test_projects/check/src/custom.erl.check b/eqwalizer/test_projects/check/src/custom.erl.check
index 9eabde9..422b0df 100644
--- a/eqwalizer/test_projects/check/src/custom.erl.check
+++ b/eqwalizer/test_projects/check/src/custom.erl.check
@@ -571,7 +571,7 @@ maps_map_2_9_neg() ->                      | ERROR   |
     M = #{a => a, b => b},                 |         |
     F = fun erlang:'=:='/2,                |         |
     maps:map(F, M).                        |         | maps:map(F, M).
-                                           |         | Expression has type:   dict  map #D{atom() => boolean()}
+                                           |         | Expression has type:   dict  map #D{'a' | 'b' => boolean()}
                                            |         | Context expected type: shape map #S{a => 'a', b => 'b'}
                                            |         |
 -spec maps_map_2_10_neg()                  |         |
@@ -3083,15 +3083,9 @@ re_replace_6_neg(Subj) ->                  | ERROR   |
     Res.                                   |         |
                                            |         |
 -spec maps_fold_keys(#{a => atom(), b => a……         |
-maps_fold_keys(M) ->                       | ERROR   |
+maps_fold_keys(M) ->                       | OK      |
     maps:fold(                             |         |
-        fun(K, _V, Acc) -> process_key(K),……         | K.
-                                           |         | Expression has type:   atom()
-                                           |         | Context expected type: 'a' | 'b'
-                                           |         | 
-                                           |         |   atom() is not compatible with 'a' | 'b'
-                                           |         |   because
-                                           |         |   atom() is not compatible with 'a'
+        fun(K, _V, Acc) -> process_key(K),……         |
         ok,                                |         |
         M                                  |         |
     ).                                     |         |
diff --git a/eqwalizer/test_projects/check/src/t_maps.erl.check b/eqwalizer/test_projects/check/src/t_maps.erl.check
index 36b158f..d4b53c8 100644
--- a/eqwalizer/test_projects/check/src/t_maps.erl.check
+++ b/eqwalizer/test_projects/check/src/t_maps.erl.check
@@ -550,15 +550,15 @@ shape_atom_key(K, Shape) ->                | OK      |
 shape_atom_key_neg(K, Shape) ->            | ERROR   |
     case Shape of                          |         |
         #{K := V} -> {K, V};               |         | {K, V}.
-                                           |         | Expression has type:   {atom(), a()}
+                                           |         | Expression has type:   {'a' | 'b', a()}
                                            |         | Context expected type: {n(), a()}
                                            |         | 
                                            |         |   at tuple index 1:
-                                           |         |   {atom(), a()} is not compatible with {n(), a()}
+                                           |         |   {'a' | 'b', a()} is not compatible with {n(), a()}
                                            |         |   because
-                                           |         |   atom() is not compatible with n()
+                                           |         |   'a' | 'b' is not compatible with n()
                                            |         |   because
-                                           |         |   atom() is not compatible with number()
+                                           |         |   'a' | 'b' is not compatible with number()
         _ -> {0, undef}                    |         |
     end.                                   |         |
                                            |         |