add some type support

src/classified.test.tsx

@@ -4,6 +4,10 @@ import { classified } from "./classified";
 
 type Props = { variant?: "primary"; block?: boolean };
 
+function Link(props: HTMLProps<HTMLAnchorElement>) {
+  return <a {...props} />;
+}
+
 test("should create a component", () => {
   let Button = classified("button")([]);
   render(<Button>hello</Button>);
@@ -23,18 +27,20 @@ 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"
   );
 });
 
 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>
@@ -47,11 +53,7 @@ test("should combine both static and dynamic classNames", () => {
 });
 
 test("should accept any react component as the first argument", () => {
-  function Link(props: HTMLProps<HTMLAnchorElement>) {
-    return <a {...props} />;
-  }
-
-  let ClassifiedLink = classified(Link)(["link"]);
+  let ClassifiedLink = classified<typeof Link | "span">(Link)(["link"]);
   render(<ClassifiedLink href="/read-more">Read more</ClassifiedLink>);
   expect(screen.getByRole("link", { name: /read more/i }).className).toBe(
     "link"
@@ -66,11 +68,21 @@ test("component should accept any additional `className`", () => {
   );
 });
 
-test("component should accept `as` prop", () => {
-  let Button = classified("button")(["btn"]);
+test("component should accept html element as `as` prop", () => {
+  let Button = classified<"button" | "a">("button")(["btn"]);
   render(
     <Button as="a" href="/">
-      hello
+      Home
+    </Button>
+  );
+  expect(screen.getByRole("link", { name: /hello/i }).className).toBe("btn");
+});
+
+test("component should accept any component as `as` prop", () => {
+  let Button = classified<"button" | typeof Link>("button")(["btn"]);
+  render(
+    <Button as={Link} href="/">
+      Home
     </Button>
   );
   expect(screen.getByRole("link", { name: /hello/i }).className).toBe("btn");

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;
   };
 }