Fixes remove styles logic (#3123)

.changeset/cold-humans-float.md

@@ -0,0 +1,5 @@
+---
+"@tokens-studio/figma-plugin": patch
+---
+
+Changed logic when "Remove styles and variables without connection to a token" is enabled where we now look at all created tokens in this session and remove them, instead of looking at each theme individually

packages/tokens-studio-for-figma/src/app/store/useTokens.tsx

@@ -438,6 +438,19 @@ export default function useTokens() {
         isInfinite: true,
       });
 
+      const allStyleIds: Record<string, string[]> = {};
+
+      const allExistingStyleReferences: string[] = themes.reduce((acc, theme) => {
+        if (theme.$figmaStyleReferences) {
+          Object.keys(theme.$figmaStyleReferences).forEach((key) => {
+            if (theme.$figmaStyleReferences && theme.$figmaStyleReferences[key]) {
+              acc.push(theme.$figmaStyleReferences[key]);
+            }
+          });
+        }
+        return acc;
+      }, [] as string[]);
+
       for (const themeId of selectedThemes) {
         const selectedTheme = themes.find((theme) => theme.id === themeId);
 
@@ -477,12 +490,30 @@ export default function useTokens() {
               selectedTheme,
             }));
 
+            Object.assign(allStyleIds, createStylesResult.styleIds);
+
             dispatch.tokenState.assignStyleIdsToCurrentTheme({ styleIds: createStylesResult.styleIds, tokens: tokensToCreate, selectedThemes });
           } else {
             notifyToUI(`No styles created for theme: ${selectedTheme.name}. Make sure some sets are enabled.`, { error: true });
           }
         }
       }
+
+      // Remove styles that aren't in the theme or in the exposed token object
+      if (settings.removeStylesAndVariablesWithoutConnection) {
+        const uniqueMergedStyleIds: string[] = Array.from(new Set([
+          ...Object.values(allExistingStyleReferences).flat(),
+          ...Object.values(allStyleIds).flat(),
+        ]));
+        const { countOfRemovedStyles } = await AsyncMessageChannel.ReactInstance.message({
+          type: AsyncMessageTypes.REMOVE_STYLES_WITHOUT_CONNECTION,
+          usedStyleIds: uniqueMergedStyleIds,
+        });
+        if (countOfRemovedStyles > 0) {
+          notifyToUI(`${countOfRemovedStyles} styles removed`);
+        }
+      }
+
       dispatch.uiState.completeJob(BackgroundJobs.UI_CREATE_STYLES);
     },
     [dispatch.tokenState, tokens, settings, themes, dispatch.uiState],

packages/tokens-studio-for-figma/src/plugin/asyncMessageHandlers/index.ts

@@ -38,3 +38,4 @@ export * from './attachLocalVariablesToTheme';
 export * from './renameVariables';
 export * from './updateVariables';
 export * from './setInitialLoad';
+export * from './removeStylesWithoutConnection';

packages/tokens-studio-for-figma/src/plugin/asyncMessageHandlers/removeStylesWithoutConnection.ts

@@ -0,0 +1,22 @@
+import { AsyncMessageChannelHandlers } from '@/AsyncMessageChannel';
+import { AsyncMessageTypes } from '@/types/AsyncMessages';
+
+export const removeStylesWithoutConnection: AsyncMessageChannelHandlers[AsyncMessageTypes.REMOVE_STYLES_WITHOUT_CONNECTION] = async (msg) => {
+  try {
+    let count = 0;
+    const localStyles = await Promise.all([
+      figma.getLocalPaintStylesAsync(),
+      figma.getLocalTextStylesAsync(),
+      figma.getLocalEffectStylesAsync(),
+    ]).then((results) => results.flat());
+    localStyles.forEach((style) => {
+      if (!msg.usedStyleIds.includes(style.id)) {
+        style.remove();
+        count += 1;
+      }
+    });
+    return { countOfRemovedStyles: count };
+  } catch (e) {
+    return { countOfRemovedStyles: 0 };
+  }
+};

packages/tokens-studio-for-figma/src/plugin/controller.ts

@@ -60,6 +60,7 @@ AsyncMessageChannel.PluginInstance.handle(
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.RESOLVE_STYLE_INFO, asyncHandlers.resolveStyleInfo);
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.SET_NONE_VALUES_ON_NODE, asyncHandlers.setNoneValuesOnNode);
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.GET_FIGMA_FONTS, asyncHandlers.getFigmaFonts);
+AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.REMOVE_STYLES_WITHOUT_CONNECTION, asyncHandlers.removeStylesWithoutConnection);
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.SET_AUTH_DATA, asyncHandlers.setAuthData);
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.CREATE_LOCAL_VARIABLES, asyncHandlers.createLocalVariables);
 AsyncMessageChannel.PluginInstance.handle(AsyncMessageTypes.CREATE_LOCAL_VARIABLES_WITHOUT_MODES, asyncHandlers.createLocalVariablesWithoutModes);

packages/tokens-studio-for-figma/src/plugin/updateStyles.test.ts

@@ -340,42 +340,4 @@ describe('updateStyles', () => {
     expect(colorSpy).not.toHaveBeenCalled();
     expect(textSpy).not.toHaveBeenCalled();
   });
-
-  it('removes styles that arent connected if setting is on', async () => {
-    const tokens = [
-      {
-        name: 'primary.500',
-        path: 'light/primary/500',
-        value: '#ff0000',
-        type: 'color',
-        styleId: '1234',
-        internal__Parent: 'global',
-      },
-    ] as ExtendedSingleToken[];
-
-    mockGetThemeInfo.mockImplementationOnce(() => (
-      Promise.resolve({
-        type: AsyncMessageTypes.GET_THEME_INFO,
-        activeTheme: {
-          [INTERNAL_THEMES_NO_GROUP]: 'light',
-        },
-        themes: [{
-          id: 'light',
-          name: 'light',
-          selectedTokenSets: {
-            global: TokenSetStatus.ENABLED,
-          },
-          $figmaStyleReferences: {
-            'primary.500': '1234',
-          },
-        }],
-      })
-    ));
-
-    await updateStyles(tokens, {
-      removeStylesAndVariablesWithoutConnection: true,
-      stylesColor: true,
-    } as SettingsState, false);
-    expect(mockRemove).toHaveBeenCalledTimes(1);
-  });
 });

packages/tokens-studio-for-figma/src/plugin/updateStyles.ts

@@ -66,21 +66,5 @@ export default async function updateStyles(
     notifyUI('Some styles were ignored due to "Ignore first part of token name" export setting', { error: true });
   }
 
-  // Remove styles that aren't in the theme or in the exposed token object
-  if (settings.removeStylesAndVariablesWithoutConnection) {
-    const [allLocalPaintStyles, allLocalTextStyles, allLocalEffectStyles] = await Promise.all([
-      figma.getLocalPaintStylesAsync(),
-      figma.getLocalTextStylesAsync(),
-      figma.getLocalEffectStylesAsync(),
-    ]);
-    const allLocalStyles = [...allLocalPaintStyles, ...allLocalTextStyles, ...allLocalEffectStyles];
-
-    allLocalStyles
-      .filter((style) => !Object.values(allStyleIds).includes(style.id))
-      .forEach((style) => {
-        style.remove();
-      });
-  }
-
   return allStyleIds;
 }

packages/tokens-studio-for-figma/src/types/AsyncMessages.ts

@@ -63,6 +63,7 @@ export enum AsyncMessageTypes {
   STARTUP = 'async/startup',
   GET_THEME_INFO = 'async/get-theme-info',
   GET_FIGMA_FONTS = 'async/get-figma-fonts',
+  REMOVE_STYLES_WITHOUT_CONNECTION = 'async/remove-styles-without-connection',
   CREATE_LOCAL_VARIABLES = 'async/create-local-variables',
   CREATE_LOCAL_VARIABLES_WITHOUT_MODES = 'async/create-local-variables-without-modes',
   RESOLVE_VARIABLE_INFO = 'async/resolve-variable-info',
@@ -261,6 +262,12 @@ export type GetFigmaFontsMessage = AsyncMessage<AsyncMessageTypes.GET_FIGMA_FONT
 export type GetFigmaFontsMessageResult = AsyncMessage<AsyncMessageTypes.GET_FIGMA_FONTS, {
   fonts: Array<Font>
 }>;
+export type RemoveStylesWithoutConnectionMessage = AsyncMessage<AsyncMessageTypes.REMOVE_STYLES_WITHOUT_CONNECTION, {
+  usedStyleIds: string[]
+}>;
+export type RemoveStylesWithoutConnectionResult = AsyncMessage<AsyncMessageTypes.REMOVE_STYLES_WITHOUT_CONNECTION, {
+  countOfRemovedStyles: number
+}>;
 export type SetAuthDataMessage = AsyncMessage<AsyncMessageTypes.SET_AUTH_DATA, {
   auth: AuthData | null
 }>;
@@ -379,7 +386,8 @@ export type AsyncMessages =
   | AttachLocalVariablesToTheme
   | RenameVariablesAsyncMessage
   | UpdateVariablesAsyncMessage
-  | RemoveRelaunchDataMessage;
+  | RemoveRelaunchDataMessage
+  | RemoveStylesWithoutConnectionMessage;
 
 export type AsyncMessageResults =
   CreateStylesAsyncMessageResult
@@ -424,7 +432,8 @@ export type AsyncMessageResults =
   | AttachLocalVariablesToThemeResult
   | RenameVariablesAsyncMessageResult
   | UpdateVariablesAsyncMessageResult
-  | RemoveRelaunchDataMessageResult;
+  | RemoveRelaunchDataMessageResult
+  | RemoveStylesWithoutConnectionResult;
 
 export type AsyncMessagesMap = {
   [K in AsyncMessageTypes]: Extract<AsyncMessages, { type: K }>