From 028628bc241ead4a3ab4c58c18a29774b5cd15a9 Mon Sep 17 00:00:00 2001
From: uhyo <uhyo@uhy.ooo>
Date: Wed, 1 May 2024 06:37:21 +0900
Subject: [PATCH] fix: improve types for .attrs() (#4288)

* fix: improve attrs() typing

* fix: update constructWithOptions type

* fix: allow double overriding

* test: add double attrs test

* fix: support double attrs pattern

* chore: remove unused generic after refactoring

---------

Co-authored-by: Evan Jacobs <ejacobs@aurorasolar.com>
---
 .../src/constructors/constructWithOptions.ts  | 119 +++++++++++-------
 packages/styled-components/src/test/types.tsx |  64 ++++++++++
 2 files changed, 137 insertions(+), 46 deletions(-)

diff --git a/packages/styled-components/src/constructors/constructWithOptions.ts b/packages/styled-components/src/constructors/constructWithOptions.ts
index e7f74fde7..c3092fc9a 100644
--- a/packages/styled-components/src/constructors/constructWithOptions.ts
+++ b/packages/styled-components/src/constructors/constructWithOptions.ts
@@ -1,7 +1,6 @@
 import {
   Attrs,
   BaseObject,
-  ExecutionProps,
   Interpolation,
   IStyledComponent,
   IStyledComponentFactory,
@@ -26,30 +25,25 @@ type AttrsResult<T extends Attrs<any>> = T extends (...args: any) => infer P
   : never;
 
 /**
- * Based on Attrs being a simple object or function that returns
- * a prop object, inspect the attrs result and attempt to extract
- * any "as" prop usage to modify the runtime target.
+ * Extract non-optional fields from given object type.
  */
-type AttrsTarget<
-  R extends Runtime,
-  T extends Attrs<any>,
-  FallbackTarget extends StyledTarget<R>,
-  Result extends ExecutionProps = AttrsResult<T>,
-> = Result extends { as: infer RuntimeTarget }
-  ? RuntimeTarget extends KnownTarget
-    ? RuntimeTarget
-    : FallbackTarget
-  : FallbackTarget;
+type RequiredFields<T, Ex> = Pick<
+  T,
+  {
+    [K in keyof T]-?: undefined extends T[K] ? never : K;
+  }[Exclude<keyof T, Ex>]
+>;
 
 export interface Styled<
   R extends Runtime,
   Target extends StyledTarget<R>,
   OuterProps extends object,
   OuterStatics extends object = BaseObject,
+  InnerProps extends object = OuterProps,
 > {
   <Props extends object = BaseObject, Statics extends object = BaseObject>(
-    initialStyles: Styles<Substitute<OuterProps, NoInfer<Props>>>,
-    ...interpolations: Interpolation<Substitute<OuterProps, NoInfer<Props>>>[]
+    initialStyles: Styles<Substitute<InnerProps, NoInfer<Props>>>,
+    ...interpolations: Interpolation<Substitute<InnerProps, NoInfer<Props>>>[]
   ): IStyledComponent<R, Substitute<OuterProps, Props>> &
     OuterStatics &
     Statics &
@@ -63,24 +57,58 @@ export interface Styled<
     Props extends object = BaseObject,
     PrivateMergedProps extends object = Substitute<OuterProps, Props>,
     PrivateAttrsArg extends Attrs<PrivateMergedProps> = Attrs<PrivateMergedProps>,
-    PrivateResolvedTarget extends StyledTarget<R> = AttrsTarget<R, PrivateAttrsArg, Target>,
   >(
     attrs: PrivateAttrsArg
-  ) => Styled<
+  ) => StyledAttrsResult<
     R,
-    PrivateResolvedTarget,
-    PrivateResolvedTarget extends KnownTarget
-      ? Substitute<
-          Substitute<OuterProps, React.ComponentPropsWithRef<PrivateResolvedTarget>>,
-          Props
-        >
-      : PrivateMergedProps,
-    OuterStatics
+    Target,
+    OuterProps,
+    OuterStatics,
+    InnerProps,
+    Props,
+    PrivateMergedProps,
+    PrivateAttrsArg
   >;
 
   withConfig: (config: StyledOptions<R, OuterProps>) => Styled<R, Target, OuterProps, OuterStatics>;
 }
 
+type StyledAttrsResult<
+  R extends Runtime,
+  Target extends StyledTarget<R>,
+  OuterProps extends object,
+  OuterStatics extends object = BaseObject,
+  InnerProps extends object = OuterProps,
+  Props extends object = BaseObject,
+  PrivateMergedProps extends object = Substitute<OuterProps, Props>,
+  PrivateAttrsArg extends Attrs<PrivateMergedProps> = Attrs<PrivateMergedProps>,
+> = (
+  AttrsResult<PrivateAttrsArg> extends { as: infer RuntimeTarget extends KnownTarget }
+    ? {
+        Target: RuntimeTarget;
+        TargetProps: Substitute<OuterProps, React.ComponentPropsWithRef<RuntimeTarget>>;
+      }
+    : { Target: Target; TargetProps: OuterProps }
+) extends {
+  Target: infer PrivateResolvedTarget extends StyledTarget<R>;
+  TargetProps: infer TargetProps extends object;
+}
+  ? Styled<
+      R,
+      PrivateResolvedTarget,
+      PrivateResolvedTarget extends KnownTarget
+        ? Substitute<TargetProps, Props & Partial<RequiredFields<PrivateAttrsArg, 'as'>>>
+        : PrivateMergedProps,
+      OuterStatics,
+      PrivateResolvedTarget extends KnownTarget
+        ? Substitute<
+            Substitute<InnerProps, React.ComponentPropsWithRef<PrivateResolvedTarget>>,
+            Props
+          >
+        : PrivateMergedProps
+    >
+  : unknown;
+
 export default function constructWithOptions<
   R extends Runtime,
   Target extends StyledTarget<R>,
@@ -88,11 +116,12 @@ export default function constructWithOptions<
     ? React.ComponentPropsWithRef<Target>
     : BaseObject,
   OuterStatics extends object = BaseObject,
+  InnerProps extends object = OuterProps,
 >(
   componentConstructor: IStyledComponentFactory<R, StyledTarget<R>, object, any>,
   tag: StyledTarget<R>,
   options: StyledOptions<R, OuterProps> = EMPTY_OBJECT
-): Styled<R, Target, OuterProps, OuterStatics> {
+): Styled<R, Target, OuterProps, OuterStatics, InnerProps> {
   /**
    * We trust that the tag is a valid component as long as it isn't
    * falsish. Typically the tag here is a string or function (i.e.
@@ -106,13 +135,13 @@ export default function constructWithOptions<
 
   /* This is callable directly as a template function */
   const templateFunction = <Props extends object = BaseObject, Statics extends object = BaseObject>(
-    initialStyles: Styles<Substitute<OuterProps, Props>>,
-    ...interpolations: Interpolation<Substitute<OuterProps, Props>>[]
+    initialStyles: Styles<Substitute<InnerProps, Props>>,
+    ...interpolations: Interpolation<Substitute<InnerProps, Props>>[]
   ) =>
-    componentConstructor<Substitute<OuterProps, Props>, Statics>(
+    componentConstructor<Substitute<InnerProps, Props>, Statics>(
       tag,
-      options as StyledOptions<R, Substitute<OuterProps, Props>>,
-      css<Substitute<OuterProps, Props>>(initialStyles, ...interpolations)
+      options as StyledOptions<R, Substitute<InnerProps, Props>>,
+      css<Substitute<InnerProps, Props>>(initialStyles, ...interpolations)
     );
 
   /**
@@ -125,24 +154,22 @@ export default function constructWithOptions<
     Props extends object = BaseObject,
     PrivateMergedProps extends object = Substitute<OuterProps, Props>,
     PrivateAttrsArg extends Attrs<PrivateMergedProps> = Attrs<PrivateMergedProps>,
-    PrivateResolvedTarget extends StyledTarget<R> = AttrsTarget<R, PrivateAttrsArg, Target>,
   >(
     attrs: PrivateAttrsArg
-  ) =>
-    constructWithOptions<
-      R,
-      PrivateResolvedTarget,
-      PrivateResolvedTarget extends KnownTarget
-        ? Substitute<
-            Substitute<OuterProps, React.ComponentPropsWithRef<PrivateResolvedTarget>>,
-            Props
-          >
-        : PrivateMergedProps,
-      OuterStatics
-    >(componentConstructor, tag, {
+  ): StyledAttrsResult<
+    R,
+    Target,
+    OuterProps,
+    OuterStatics,
+    InnerProps,
+    Props,
+    PrivateMergedProps,
+    PrivateAttrsArg
+  > =>
+    constructWithOptions<R, Target, any, any, any>(componentConstructor, tag, {
       ...options,
       attrs: Array.prototype.concat(options.attrs, attrs).filter(Boolean),
-    });
+    }) as any;
 
   /**
    * If config methods are called, wrap up a new template function
diff --git a/packages/styled-components/src/test/types.tsx b/packages/styled-components/src/test/types.tsx
index a9fd46cc3..28c98bfc2 100644
--- a/packages/styled-components/src/test/types.tsx
+++ b/packages/styled-components/src/test/types.tsx
@@ -270,6 +270,70 @@ const AttrRequiredTest4 = styled(DivWithUnfulfilledRequiredProps).attrs({
   waz: 42,
 })``;
 
+{
+  const DivWithRequiredFooBar = styled.div<{ foo: number; bar: string }>``;
+  // @ts-expect-error must provide both foo and bar
+  <DivWithRequiredFooBar />;
+  // @ts-expect-error must provide both foo and bar
+  <DivWithRequiredFooBar foo={3} />;
+  // @ts-expect-error must provide both foo and bar
+  <DivWithRequiredFooBar bar="3" />;
+  // OK
+  <DivWithRequiredFooBar foo={3} bar="3" />;
+
+  // foo is provided, so it becomes optional
+  const DivWithRequiredBar = styled(DivWithRequiredFooBar).attrs({ foo: 42 })`
+    margin; ${props => props.foo * 10}px;
+  `;
+  // @ts-expect-error must provide bar
+  <DivWithRequiredBar />;
+  // OK
+  <DivWithRequiredBar bar="3" />;
+  // OK. Can still provide foo if we want
+  <DivWithRequiredBar foo={3} bar="3" />;
+  // @ts-expect-error foo must be a number
+  <DivWithRequiredBar foo="3" bar="3" />;
+
+  const Div = styled(DivWithRequiredBar).attrs({ bar: '42' })`
+    margin: ${props => {
+      // @ts-expect-error foo is optional
+      const foo: number = props.foo;
+      const bar: string = props.bar;
+      return foo * Number(bar);
+    }}px;
+  `;
+  // OK
+  <Div />;
+  <Div foo={3} />;
+  <Div bar="3" />;
+  <Div foo={3} bar="3" />;
+}
+
+{
+  // double attrs
+  const DivWithRequiredFooBar = styled.div<{ foo: number; bar: number }>``;
+  const Div = styled(DivWithRequiredFooBar).attrs({ foo: 42 }).attrs({ bar: 42 })`
+    margin: ${props => props.foo * props.bar}px;
+    `;
+
+  <Div />;
+  <Div foo={3} />;
+  <Div bar={3} />;
+  <Div foo={3} bar={3} />;
+}
+
+{
+  // double overriding
+  const Div = styled.div``;
+  const H1 = styled(Div).attrs({ as: 'h1' })``;
+  const Label = styled(H1).attrs({ as: 'label' })``;
+
+  <Label
+    ref={(el: HTMLLabelElement | null) => {}}
+    onCopy={(e: React.ClipboardEvent<HTMLLabelElement>) => {}}
+  />;
+}
+
 /** Intrinsic props and ref are being incorrectly types when using `as`
  * https://github.com/styled-components/styled-components/issues/3800#issuecomment-1548941843
  */