From 374aa99d32da61fd4ff0a192e3a59a64fe041966 Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Mon, 13 Apr 2020 08:37:17 +0800
Subject: [PATCH] Improve typings for Source enhancement

---
 addons/docs/src/blocks/enhanceSource.ts | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/addons/docs/src/blocks/enhanceSource.ts b/addons/docs/src/blocks/enhanceSource.ts
index 8f520631368..1d3cf583d61 100644
--- a/addons/docs/src/blocks/enhanceSource.ts
+++ b/addons/docs/src/blocks/enhanceSource.ts
@@ -1,4 +1,5 @@
-import { ArgsEnhancer, combineParameters } from '@storybook/client-api';
+import { combineParameters } from '@storybook/client-api';
+import { StoryContext, Parameters } from '@storybook/addons';
 
 interface Location {
   line: number;
@@ -36,7 +37,7 @@ const extract = (targetId: string, { source, locationsMap }: StorySource) => {
   ].join('\n');
 };
 
-export const enhanceSource: ArgsEnhancer = (context) => {
+export const enhanceSource = (context: StoryContext): Parameters => {
   const { id, parameters } = context;
   const { storySource, docs = {} } = parameters;
   const { formatSource } = docs;