refactor, clean up, fix some other cases

code/renderers/react/src/__test-dts__/CSF3.test-d.tsx

@@ -85,9 +85,8 @@ describe('Story args can be inferred', () => {
 
     const meta = satisfies<Meta<Props>>()({
       component: Button,
-      args: { label: 'good', disabled: false },
+      args: { disabled: false },
       render: (args, { component }) => {
-        // TODO: Might be nice if we can infer that.
         // component is not null as it is provided in meta
         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
         const Component = component!;
@@ -99,28 +98,53 @@ describe('Story args can be inferred', () => {
       },
     });
 
-    const Basic: StoryObj<typeof meta> = { args: { theme: 'light' } };
+    const Basic: StoryObj<typeof meta> = { args: { theme: 'light', label: 'good' } };
   });
 
+  const withDecorator: DecoratorFn<{ decoratorArg: number }> = (Story, { args }) => (
+    <>
+      Decorator: {args.decoratorArg}
+      This Story allows optional TArgs, but the decorator only knows about the decoratorArg. It
+      should really allow optionally a Partial of TArgs.
+      <Story args={{ decoratorArg: 0 }} />
+    </>
+  );
+
   test('Correct args are inferred when type is widened for decorators', () => {
     type Props = ButtonProps & { decoratorArg: number };
 
-    const withDecorator: DecoratorFn<{ decoratorArg: number }> = (Story, { args }) => (
-      <>
-        Decorator: {args.decoratorArg}
-        This Story allows optional TArgs, but the decorator only knows about the decoratorArg. It
-        should really allow optionally a Partial of TArgs.
-        <Story args={{ decoratorArg: 0 }} />
-      </>
-    );
-
     const meta = satisfies<Meta<Props>>()({
       component: Button,
-      args: { label: 'good', disabled: false },
+      args: { disabled: false },
       decorators: [withDecorator],
     });
 
     // Yes, decorator arg is required
-    const Basic: StoryObj<typeof meta> = { args: { decoratorArg: 0 } };
+    const Basic: StoryObj<typeof meta> = { args: { decoratorArg: 0, label: 'good' } };
+  });
+
+  test('Correct args are inferred when type is widened for decorators and render functions', () => {
+    type Props = ButtonProps & { decoratorArg: number } & { theme: ThemeData };
+
+    const meta = satisfies<Meta<Props>>()({
+      component: Button,
+      args: { disabled: false },
+      decorators: [withDecorator],
+      loaders: [async ({ args }: { args: { label: string } }) => ({ data: args.label })],
+      render: (args, { component }) => {
+        // component is not null as it is provided in meta
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        const Component = component!;
+        return (
+          <Theme theme={args.theme}>
+            <Component {...args} />
+          </Theme>
+        );
+      },
+    });
+
+    const Basic: StoryObj<typeof meta> = {
+      args: { decoratorArg: 0, label: 'good', theme: 'light' },
+    };
   });
 });

code/renderers/react/src/public-types.ts

@@ -1,7 +1,14 @@
-import { AnnotatedStoryFn, Args, ComponentAnnotations, StoryAnnotations } from '@storybook/csf';
+import type {
+  AnnotatedStoryFn,
+  Args,
+  ComponentAnnotations,
+  StoryAnnotations,
+  ArgsStoryFn,
+  DecoratorFunction,
+  LoaderFunction,
+} from '@storybook/csf';
 import { ComponentType, JSXElementConstructor } from 'react';
 import { ReactFramework } from './types';
-import { ArgsStoryFn, DecoratorFunction } from '../../../../../csf/src';
 
 type JSXElement = keyof JSX.IntrinsicElements | JSXElementConstructor<any>;
 
@@ -10,12 +17,9 @@ type JSXElement = keyof JSX.IntrinsicElements | JSXElementConstructor<any>;
  *
  * @see [Default export](https://storybook.js.org/docs/formats/component-story-format/#default-export)
  */
-export type Meta<
-  CmpOrArgs = Args,
-  StoryArgs = CmpOrArgs extends ComponentType<infer CmpArgs> ? CmpArgs : CmpOrArgs
-> = CmpOrArgs extends ComponentType<infer CmpArgs>
-  ? ComponentAnnotations<ReactFramework<CmpArgs>, StoryArgs>
-  : ComponentAnnotations<ReactFramework<CmpOrArgs>, StoryArgs>;
+export type Meta<CmpOrArgs = Args> = CmpOrArgs extends ComponentType<infer CmpArgs>
+  ? ComponentAnnotations<ReactFramework, CmpArgs>
+  : ComponentAnnotations<ReactFramework, CmpOrArgs>;
 
 /**
  * Story function that represents a CSFv2 component example.
@@ -29,23 +33,29 @@ export type StoryFn<TArgs = Args> = AnnotatedStoryFn<ReactFramework, TArgs>;
  *
  * @see [Named Story exports](https://storybook.js.org/docs/formats/component-story-format/#named-story-exports)
  */
+
 export type StoryObj<MetaOrArgs = Args> = MetaOrArgs extends {
-  render?: ArgsStoryFn<ReactFramework, infer StoryArgs>;
-  decorators?: DecoratorFunction<ReactFramework, infer StoryArgs>[];
+  render?: ArgsStoryFn<ReactFramework, infer RArgs>;
+  decorators?: DecoratorFunction<ReactFramework, infer DArgs>[];
   component?: ComponentType<infer CmpArgs>;
-  args?: infer D;
+  loaders?: LoaderFunction<ReactFramework, infer LArgs>[];
+  args?: infer DefaultArgs;
 }
-  ? (unknown extends StoryArgs ? CmpArgs : StoryArgs) extends infer Args
-    ? StoryAnnotations<ReactFramework<CmpArgs>, Args> & StrictStoryArgs<Args, D>
+  ? CmpArgs & RArgs & DArgs & LArgs extends infer Args
+    ? StoryAnnotations<
+        ReactFramework,
+        Args,
+        SetOptional<Args, keyof (DefaultArgs & ActionArgs<Args>)>
+      >
     : never
   : StoryAnnotations<ReactFramework, MetaOrArgs>;
 
-type StrictStoryArgs<Args, D> = {} extends MakeOptional<Args, D & ActionArgs<Args>>
-  ? { args?: Partial<Args> }
-  : { args: MakeOptional<Args, D & ActionArgs<Args>> };
-
 type ActionArgs<Args> = {
   [P in keyof Args as ((...args: any[]) => void) extends Args[P] ? P : never]: Args[P];
 };
 
-type MakeOptional<T, O> = Omit<T, keyof O> & Partial<Pick<T, Extract<keyof T, keyof O>>>;
+type SetOptional<T, K> = {
+  [P in keyof T as P extends K ? P : never]?: T[P];
+} & {
+  [P in keyof T as P extends K ? never : P]: T[P];
+};

code/renderers/react/src/types.ts

@@ -1,12 +1,13 @@
 import type { ComponentType, ReactElement } from 'react';
+import type { AnyFramework } from '@storybook/csf';
 
 export type { RenderContext } from '@storybook/store';
 export type { StoryContext } from '@storybook/csf';
 
-export type ReactFramework<CmpArgs = any> = {
-  component: ComponentType<CmpArgs>;
-  storyResult: StoryFnReactReturnType;
-};
+// export interface ReactFramework extends AnyFramework {
+//   component: ComponentType<this['T']>;
+//   storyResult: StoryFnReactReturnType;
+// }
 
 export interface ShowErrorArgs {
   title: string;
@@ -24,3 +25,23 @@ export interface IStorybookSection {
   kind: string;
   stories: IStorybookStory[];
 }
+
+export type Framework = {
+  component: unknown;
+  T: unknown; // higher kinded type
+  storyResult: unknown;
+};
+
+export interface ReactFramework extends Framework {
+  component: ComponentType<this['T']>; // reference the higher kinded type
+  storyResult: StoryFnReactReturnType;
+}
+
+// used like this:
+interface ComponentAnnotations<TFramework extends Framework, TArgs> {
+  // specify the generic type with an intersection
+  component?: (TFramework & { T: TArgs })['component'];
+
+  // falls back to unknown, if you don't intersect T
+  subcomponents?: Record<string, TFramework['component']>;
+}