Repository: jeremydaly/dynamodb-toolbox (ts)

Issue:
Cannot use 'set' type Entity attribute with explicit model type
Using a generic specifier on an Entity with an attribute of e.g. `{ type: 'set', setType: 'string' }` causes the following TS error on update: 
> Type '{ $delete: string[]; }' is not assignable to type 'string[]'.

The same issue occurs with typing a number and using $add to increase it during an update.

Example:
```ts
const Customer = new Entity<{
  id: string;
  sk: string;
  count: number;
  stringSet: string[];
}>({
  name: 'Customer',
  attributes: {
    id: { partitionKey: true }, // flag as partitionKey
    sk: { sortKey: true }, // flag as sortKey
    count: { type: 'number' },
    stringSet: { type: 'set', setType: 'string' }
  },
  table: MyTable
});

Customer.update({
  id: 'foo',
  sk: 'bar',
  count: { $add: 1 }, // TS error similar to the above 
  stringSet: { $add: ['baz'] } // above TS error
});
```

It works without the generic specifier on Entity.

Pull request:
feat(entity/schema/sets): enhance typings of set types attrs
closes #141

Patch:
diff --git a/src/classes/Entity/types.ts b/src/classes/Entity/types.ts
index 1b893053..814fb33d 100644
--- a/src/classes/Entity/types.ts
+++ b/src/classes/Entity/types.ts
@@ -2,7 +2,7 @@ import type { DocumentClient } from 'aws-sdk/clients/dynamodb'
 import type { A, B, O, F } from 'ts-toolbelt'
 
 import type { FirstDefined, If } from '../../lib/utils'
-import type { DynamoDBKeyTypes, DynamoDBTypes, $QueryOptions, TableDef } from '../Table'
+import type { DynamoDBKeyTypes, DynamoDBTypes, $QueryOptions, TableDef } from '../Table';
 import Entity from './Entity'
 
 export interface EntityConstructor<
@@ -227,7 +227,8 @@ export type InferItemAttributeValue<
     | GSISortKeyDefinition
     | PureAttributeDefinition
     ? Definition['type'] extends DynamoDBTypes
-      ? FromDynamoData<A.Cast<Definition['type'], DynamoDBTypes>>
+      ? Definition['setType'] extends DynamoDBKeyTypes
+       ? FromDynamoData<NonNullable<Definition['setType']>>[] : FromDynamoData<A.Cast<Definition['type'], DynamoDBTypes>>
       : any
     : never
   composite: Definition extends CompositeAttributeDefinition
@@ -437,39 +438,34 @@ export interface UpdateCustomParameters {
 
 export type UpdateCustomParams = O.Partial<UpdateCustomParameters & DocumentClient.UpdateItemInput>
 
-export type UpdateItem<
-  MethodItemOverlay extends Overlay,
+export type UpdateItem<MethodItemOverlay extends Overlay,
   EntityItemOverlay extends Overlay,
   CompositePrimaryKey extends O.Object,
   Item extends O.Object,
   Attributes extends ParsedAttributes,
-  StrictSchemaCheck extends boolean | undefined = true
-> = FirstDefined<
-  [
+  StrictSchemaCheck extends boolean | undefined = true> = FirstDefined<[
     MethodItemOverlay,
     EntityItemOverlay,
-    A.Compute<
-      CompositePrimaryKey &
-        {
-          [inputAttr in Attributes['always']['input']]:
-            | Item[A.Cast<inputAttr, keyof Item>]
-            | { $delete?: string[]; $add?: any; $prepend?: any[]; $append?: any[] }
-        } &
-        {
-          [optAttr in Attributes['required']['all'] | Attributes['always']['default']]?:
-            | Item[A.Cast<optAttr, keyof Item>]
-            | { $delete?: string[]; $add?: any; $prepend?: any[]; $append?: any[] }
-        } &
-        {
-          [attr in Attributes['optional']]?:
-            | null
-            | Item[A.Cast<attr, keyof Item>]
-            | { $delete?: string[]; $add?: any; $append?: any[]; $prepend?: any[] }
-        } & { $remove?: Attributes['optional'] | Attributes['optional'][] }
-    >
+    A.Compute<CompositePrimaryKey &
+      {
+        [inputAttr in Attributes['always']['input'] & keyof Item]: AttributeUpdateInput<Item[inputAttr]>
+      } &
+      {
+        [inputRequiredOrWithDefaultAttribute in (Attributes['required']['all'] | Attributes['always']['default']) & keyof Item]?: AttributeUpdateInput<Item[inputRequiredOrWithDefaultAttribute]>
+      } &
+      {
+        [inputOptionalAttribute in Attributes['optional'] & keyof Item]?: AttributeUpdateInput<Item[inputOptionalAttribute]> | null
+      } & { $remove?: Attributes['optional'] | Attributes['optional'][] }>
   ]
-  | If<A.Equals<StrictSchemaCheck, true>, never, any>
->
+  | If<A.Equals<StrictSchemaCheck, true>, never, any>>
+
+export type AttributeUpdateInput<AttributeType> =
+    | If<A.Equals<AttributeType, FromDynamoData<'list' | 'set'> | undefined>, { $delete?: string[]; $add?: any; $prepend?: AttributeType; $append?: AttributeType; $remove?: number[] } | AttributeType, AttributeType>
+    | If<A.Equals<AttributeType, number[] | undefined>, { $delete?: number[]; $add?: number[]; $prepend?: AttributeType; $append?: number[]; } | string[]>
+    | If<A.Equals<AttributeType, string[] | undefined>, { $delete?: string[]; $add?: string[]; $prepend?: AttributeType; $append?: string[]; } | number[]>
+    | If<A.Equals<AttributeType, boolean[] | undefined>, { $delete?: boolean[]; $add?: boolean[]; $prepend?: AttributeType; $append?: boolean[]; } | boolean[]>
+    | If<A.Equals<AttributeType, FromDynamoData<'number'> | undefined>, { $add?: number }>
+
 
 export type DeleteOptionsReturnValues = 'NONE' | 'ALL_OLD'
 


Test patch:
diff --git a/src/__tests__/entity.put.unit.test.ts b/src/__tests__/entity.put.unit.test.ts
index f6567015..f8ac5f95 100644
--- a/src/__tests__/entity.put.unit.test.ts
+++ b/src/__tests__/entity.put.unit.test.ts
@@ -354,7 +354,7 @@ describe('put', () => {
       TestEntity.putParams({
         email: 'test-pk',
         sort: 'test-sk',
-        // 💥 TODO: Improve set support
+        // @ts-expect-error
         test_string_set_type: [1, 2, 3]
       })
     ).toThrow(`'test_string_set_type' must be a valid set (array) containing only string types`)
diff --git a/src/__tests__/entity.update.unit.test.ts b/src/__tests__/entity.update.unit.test.ts
index f6379110..a95e86cc 100644
--- a/src/__tests__/entity.update.unit.test.ts
+++ b/src/__tests__/entity.update.unit.test.ts
@@ -409,10 +409,9 @@ describe('update', () => {
       ExpressionAttributeNames,
       ExpressionAttributeValues
     } = TestEntity.updateParams({
-      pk: 'test-pk',
-      sk: 'test-sk',
+      email: 'test-pk',
+      sort: 'test-sk',
       test_string_set_type: { $delete: ['1', '2', '3'] },
-      // @ts-expect-error 💥 TODO: Improve list support
       test_number_set_type: { $delete: [1, 2, 3] }
     })
     expect(UpdateExpression).toBe(
@@ -450,7 +449,6 @@ describe('update', () => {
     } = TestEntity.updateParams({
       email: 'test-pk',
       sort: 'test-sk',
-      // @ts-expect-error 💥 TODO: Improve list support
       test_list: { $remove: [2, 3, 8] }
     })
     expect(UpdateExpression).toBe(
@@ -798,8 +796,8 @@ describe('update', () => {
   it('fails when using non-numeric values for indexed list removals', () => {
     expect(() =>
       TestEntity.updateParams({
-        pk: 'test-pk',
-        sk: 'test-sk',
+        email: 'test-pk',
+        sort: 'test-sk',
         // @ts-expect-error
         test_list: { $remove: [1, 2, 'test'] }
       })