diff --git a/index.d.ts b/index.d.ts
index bed6df26..22906002 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -3914,51 +3914,53 @@ type PartialDefinedKeys<T> = {
 };
 
 export type ItemAttribute<A extends Attribute> =
-  A["type"] extends OpaquePrimitiveTypeName<infer T>
-    ? T
-    : A["type"] extends CustomAttributeTypeName<infer T>
-    ? T
-    : A["type"] extends infer R
-    ? R extends "string"
-      ? string
-      : R extends "number"
-      ? number
-      : R extends "boolean"
-      ? boolean
-      : R extends ReadonlyArray<infer E>
-      ? E
-      : R extends "map"
-      ? "properties" extends keyof A
-        ? {
-            [P in keyof A["properties"]]: A["properties"][P] extends infer M
-              ? M extends Attribute
-                ? ItemAttribute<M>
-                : never
-              : never;
-          }
-        : never
-      : R extends "list"
-      ? "items" extends keyof A
-        ? A["items"] extends infer I
-          ? I extends Attribute
-            ? Array<ItemAttribute<I>>
+  A["type"] extends infer T
+    ? T extends OpaquePrimitiveTypeName<infer OP>
+      ? OP
+      : T extends CustomAttributeTypeName<infer CA>
+      ? CA
+      : T extends infer R
+      ? R extends "string"
+        ? string
+        : R extends "number"
+        ? number
+        : R extends "boolean"
+        ? boolean
+        : R extends ReadonlyArray<infer E>
+        ? E
+        : R extends "map"
+        ? "properties" extends keyof A
+          ? {
+              [P in keyof A["properties"]]: A["properties"][P] extends infer M
+                ? M extends Attribute
+                  ? ItemAttribute<M>
+                  : never
+                : never;
+            }
+          : never
+        : R extends "list"
+        ? "items" extends keyof A
+          ? A["items"] extends infer I
+            ? I extends Attribute
+              ? Array<ItemAttribute<I>>
+              : never
             : never
           : never
-        : never
-      : R extends "set"
-      ? "items" extends keyof A
-        ? A["items"] extends infer I
-          ? I extends "string"
-            ? string[]
-            : I extends "number"
-            ? number[]
-            : I extends ReadonlyArray<infer ENUM>
-            ? ENUM[]
+        : R extends "set"
+        ? "items" extends keyof A
+          ? A["items"] extends infer I
+            ? I extends "string"
+              ? string[]
+              : I extends "number"
+              ? number[]
+              : I extends ReadonlyArray<infer ENUM>
+              ? ENUM[]
+              : never
             : never
           : never
+        : R extends "any"
+        ? any
         : never
-      : R extends "any"
-      ? any
       : never
     : never;
 
diff --git a/test/entity.test-d.ts b/test/entity.test-d.ts
index fcb09be4..ba2920fa 100644
--- a/test/entity.test-d.ts
+++ b/test/entity.test-d.ts
@@ -1,4 +1,4 @@
-import { Entity, Resolve } from "../";
+import { Entity, Resolve, CustomAttributeType, EntityRecord } from "../";
 import { expectType } from "tsd";
 
 const troubleshoot = <Params extends any[], Response>(
@@ -151,6 +151,42 @@ const requiredMapAttributeEntity = new Entity({
   },
 });
 
+type UnionType =
+  | { prop1: string }
+  | { prop1: string; prop2: number }
+  | { prop3: string }
+  | { prop4: number; prop3: string };
+
+const customAttributeEntity = new Entity({
+  model: {
+    entity: "user",
+    service: "versioncontrol",
+    version: "1",
+  },
+  attributes: {
+    union: {
+      required: true,
+      type: CustomAttributeType<UnionType>("any"),
+    },
+    stringVal: {
+      type: "string",
+    },
+  },
+  indexes: {
+    user: {
+      collection: "overview",
+      pk: {
+        composite: ["stringVal"],
+        field: "pk",
+      },
+      sk: {
+        composite: ["stringVal2"],
+        field: "sk",
+      },
+    },
+  },
+});
+
 entityWithSK
   .update({
     attr1: "abc",
@@ -200,3 +236,8 @@ expectType<{
     test?: string;
   };
 }>(magnify(item));
+
+type CustomAttributeEntityType = EntityRecord<typeof customAttributeEntity>;
+const unionItem = {} as CustomAttributeEntityType["union"];
+
+expectType<UnionType>(magnify(unionItem));