From 5583f19d7504a2539798236a042afc81beee313d Mon Sep 17 00:00:00 2001
From: VKotarac <vladan.kotarac@odt.net>
Date: Tue, 24 Jan 2023 11:42:01 +0100
Subject: [PATCH] add some type support

---
 src/classified.test.tsx | 10 ++++++----
 src/classified.ts       | 33 +++++++++++++++++++++++----------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/src/classified.test.tsx b/src/classified.test.tsx
index 4eb24b0..eb52d98 100644
--- a/src/classified.test.tsx
+++ b/src/classified.test.tsx
@@ -23,7 +23,9 @@ test("should apply dynamic classNames", () => {
 });
 
 test("should pass props to dynamic className creator", () => {
-  let Button = classified("button")([({ variant }: Props) => `btn-${variant}`]);
+  let Button = classified<"button", Props>("button")([
+    ({ variant }) => `btn-${variant}`,
+  ]);
   render(<Button variant="primary">hello</Button>);
   expect(screen.getByRole("button", { name: /hello/i }).className).toBe(
     "btn-primary"
@@ -31,10 +33,10 @@ test("should pass props to dynamic className creator", () => {
 });
 
 test("should combine both static and dynamic classNames", () => {
-  let Button = classified("button")([
+  let Button = classified<"button", Props>("button")([
     "btn",
-    ({ variant }: Props) => `btn-${variant}`,
-    ({ block }: Props) => block && "btn-block",
+    ({ variant }) => `btn-${variant}`,
+    ({ block }) => block && "btn-block",
   ]);
   render(
     <Button variant="primary" block>
diff --git a/src/classified.ts b/src/classified.ts
index 4ed8270..e9c7ab7 100644
--- a/src/classified.ts
+++ b/src/classified.ts
@@ -1,28 +1,41 @@
-import type { ComponentProps, ElementType, HTMLProps } from "react";
+import type { ComponentProps, ElementType } from "react";
 import { createElement, forwardRef } from "react";
 import isPropValid from "@emotion/is-prop-valid";
 
 type Falsy = false | 0 | "" | null | undefined;
 
-function classified(type) {
-  return function createClassifiedComponent(names) {
-    function ClassifiedComponent({ as = type, children, ...props }, ref) {
+type Generator<Props> = ((props: Props) => string | Falsy) | string | Falsy;
+
+type ClassNames<Props> = Array<Generator<Props>>;
+
+function classified<
+  Type extends ElementType,
+  AdditionalProps extends object = {}
+>(type: Type) {
+  type GeneratorProps = AdditionalProps & ComponentProps<Type>;
+
+  return function createClassifiedComponent(names: ClassNames<GeneratorProps>) {
+    type Props = GeneratorProps & {
+      as?: Type;
+      className?: string;
+    };
+
+    const ClassifiedComponent = forwardRef<Type, Props>((props: Props, ref) => {
       let className = names
         .concat(props.className)
         .map(name => (isFunc(name) ? name(props) : name))
         .filter(Boolean)
         .join(" ");
 
-      let elementProps = isString(as) ? getHtmlProps(props) : props;
+      let elementProps = isString(props.as) ? getHtmlProps(props) : props;
 
       return createElement(
-        as,
-        Object.assign(elementProps, { className, ref }),
-        children
+        props.as || type,
+        Object.assign({}, elementProps, { className, ref })
       );
-    }
+    });
 
-    return forwardRef(ClassifiedComponent);
+    return ClassifiedComponent;
   };
 }