From ade8805296f0ce629393651734ffd4c0ea9096ca Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@chromatic.com>
Date: Mon, 6 Mar 2023 09:49:17 +1100
Subject: [PATCH] Fix types of `useOf`

---
 code/ui/blocks/src/blocks/Canvas.tsx | 4 ++--
 code/ui/blocks/src/blocks/useOf.ts   | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/code/ui/blocks/src/blocks/Canvas.tsx b/code/ui/blocks/src/blocks/Canvas.tsx
index 77a2fbc6fff2..d354fe775ad3 100644
--- a/code/ui/blocks/src/blocks/Canvas.tsx
+++ b/code/ui/blocks/src/blocks/Canvas.tsx
@@ -1,7 +1,7 @@
 /* eslint-disable react/destructuring-assignment */
 import React, { Children, useContext } from 'react';
 import type { FC, ReactElement, ReactNode } from 'react';
-import type { ModuleExport, ModuleExports, Renderer } from '@storybook/types';
+import type { ModuleExport, ModuleExports, PreparedStory, Renderer } from '@storybook/types';
 import { deprecate } from '@storybook/client-logger';
 import dedent from 'ts-dedent';
 import type { Layout, PreviewProps as PurePreviewProps } from '../components';
@@ -157,7 +157,7 @@ export const Canvas: FC<CanvasProps & DeprecatedCanvasProps> = (props) => {
   const { children, of, source } = props;
   const { isLoading, previewProps } = useDeprecatedPreviewProps(props, docsContext, sourceContext);
 
-  let story;
+  let story: PreparedStory;
   let sourceProps;
   /**
    * useOf and useSourceProps will throw if they can't find the story, in the scenario where
diff --git a/code/ui/blocks/src/blocks/useOf.ts b/code/ui/blocks/src/blocks/useOf.ts
index cde765f061ac..8e5d9d5f7a4f 100644
--- a/code/ui/blocks/src/blocks/useOf.ts
+++ b/code/ui/blocks/src/blocks/useOf.ts
@@ -1,4 +1,9 @@
-import type { DocsContextProps, ModuleExport, ResolvedModuleExportType } from '@storybook/types';
+import type {
+  DocsContextProps,
+  ModuleExport,
+  ResolvedModuleExportType,
+  ResolvedModuleExportFromType,
+} from '@storybook/types';
 import { useContext } from 'react';
 import { DocsContext } from './DocsContext';
 
@@ -13,7 +18,7 @@ export type Of = Parameters<DocsContextProps['resolveOf']>[0];
 export const useOf = <TType extends ResolvedModuleExportType>(
   moduleExportOrType: ModuleExport | TType,
   validTypes?: TType[]
-): ReturnType<DocsContextProps['resolveOf']> => {
+): ResolvedModuleExportFromType<TType> => {
   const context = useContext(DocsContext);
   return context.resolveOf(moduleExportOrType, validTypes);
 };