feat: enhance listdata selection methods (#8656)

* add addSelectedKeys and removeSelectedKeys  methods to useListData

* create tests for new methods to useListData

* keep same pattern for reassigning selection variable

* rename methods

* keep original selection value

* remove validation for keys

* adding some additional test cases

---------

Co-authored-by: rbuzatto <[email protected]>
Co-authored-by: Daniel Lu <[email protected]>

Author: 3 people

packages/@react-stately/data/src/useListData.ts

@@ -36,6 +36,12 @@ export interface ListData<T> {
   /** Sets the selected keys. */
   setSelectedKeys(keys: Selection): void,
 
+  /** Adds the given keys to the current selected keys. */
+  addKeysToSelection(keys: Selection): void,
+
+  /** Removes the given keys from the current selected keys. */
+  removeKeysFromSelection(keys: Selection): void,
+
   /** The current filter text. */
   filterText: string,
 
@@ -175,6 +181,43 @@ export function createListActions<T, C>(opts: CreateListOptions<T, C>, dispatch:
         selectedKeys
       }));
     },
+    addKeysToSelection(selectedKeys: Selection) {
+      dispatch(state => {
+        if (state.selectedKeys === 'all') {
+          return state;
+        }
+        if (selectedKeys === 'all') {
+          return {
+            ...state,
+            selectedKeys: 'all'
+          };
+        }
+
+        return {
+          ...state,
+          selectedKeys: new Set([...state.selectedKeys, ...selectedKeys])
+        };
+      });
+    },
+    removeKeysFromSelection(selectedKeys: Selection) {
+      dispatch(state => {
+        if (selectedKeys === 'all') {
+          return {
+            ...state,
+            selectedKeys: new Set()
+          };
+        }
+
+        let selection: Selection = state.selectedKeys  === 'all' ? new Set(state.items.map(getKey!)) : new Set(state.selectedKeys);
+        for (let key of selectedKeys) {
+          selection.delete(key);
+        }
+        return {
+          ...state,
+          selectedKeys: selection
+        };
+      });
+    },
     setFilterText(filterText: string) {
       dispatch(state => ({
         ...state,

packages/@react-stately/data/test/useListData.test.js

@@ -44,6 +44,85 @@ describe('useListData', function () {
     expect(result.current.selectedKeys).toEqual(new Set(['Sam', 'Julia']));
   });
 
+  describe('addKeysToSelection', function () {
+    it('should add selected keys', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam']}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.addKeysToSelection(['Julia']);
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual(new Set(['Sam', 'Julia']));
+    });
+
+    it('should support adding "all" to selected keys', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam']}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.addKeysToSelection('all');
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual('all');
+    });
+
+    it('should still return "all" if selected keys was already "all"', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: 'all'}));
+
+      act(() => {
+        result.current.addKeysToSelection(['Same']);
+      });
+      expect(result.current.selectedKeys).toEqual('all');
+    });
+  });
+
+  describe('removeKeysFromSelection', function () {
+    it('should remove all keys', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam', 'Julia']}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.removeKeysFromSelection('all');
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual(new Set());
+    });
+
+    it('should remove the selected keys', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam', 'Julia']}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.removeKeysFromSelection(['Sam']);
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual(new Set(['Julia']));
+    });
+
+    it('should remove the selected keys from an "all" set', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: 'all'}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.removeKeysFromSelection(['Sam', 'David']);
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual(new Set(['Julia']));
+    });
+
+    it('should support removing "all"', function () {
+      let {result} = renderHook(() => useListData({initialItems: initial, getKey, initialSelectedKeys: ['Sam', 'Julia']}));
+      let initialResult = result.current;
+
+      act(() => {
+        result.current.removeKeysFromSelection('all');
+      });
+      expect(result.current.selectedKeys).not.toBe(initialResult.selectedKeys);
+      expect(result.current.selectedKeys).toEqual(new Set([]));
+    });
+  });
+
   it('should get an item by key', function () {
     let {result} = renderHook(() => useListData({initialItems: initial, getKey}));
     expect(result.current.getItem('Sam')).toBe(initial[1]);