Add forward refs.

Author: aidanCQ

src/custom/theme-selector.tsx

@@ -13,7 +13,7 @@ export const useTheme = () => {
   return { theme, setMode: (_mode: typeof theme['mode']) => setTheme(_mode) }
 }
 
-export const ThemeSelector = ({theme, setMode}: ReturnType<typeof useTheme>) => {
+export const ThemeSelector = React.forwardRef<HTMLButtonElement, ReturnType<typeof useTheme>>(({theme, setMode}, ref) => {
 
   const stateMap = {
     "light": {
@@ -31,9 +31,9 @@ export const ThemeSelector = ({theme, setMode}: ReturnType<typeof useTheme>) =>
       if (theme.mode === "dark") setMode('light')
       if (theme.mode === "light") setMode("system")
       if (theme.mode === "system") setMode("dark")
-    }}>
+    }} ref={ref}>
       {stateMap[theme.mode].icon}
     </Button>
   )
- 
-};
+
+})

src/shadcn/ui/badge.tsx

@@ -27,10 +27,10 @@ export interface BadgeProps
   extends React.HTMLAttributes<HTMLDivElement>,
     VariantProps<typeof badgeVariants> {}
 
-function Badge({ className, variant, ...props }: BadgeProps) {
+const Badge = React.forwardRef<HTMLDivElement, React.ComponentProps<'div'> & BadgeProps>(({ className, variant, ...props }, ref )  => {
   return (
-    <div className={cn(badgeVariants({ variant }), className)} {...props} />
+    <div className={cn(badgeVariants({ variant }), className)} {...props}  ref={ref}/>
   )
-}
+})
 
 export { Badge, badgeVariants }

src/shadcn/ui/skeleton.tsx

@@ -1,20 +1,16 @@
+import React from "react"
 import { cn } from "src/utils"
 
-function Skeleton({
-  className,
-  ...props
-}: React.HTMLAttributes<HTMLDivElement>) {
-  return (
-    <div
+const Skeleton = React.forwardRef<HTMLDivElement, React.ComponentProps<"div">>(({className, ...props}, ref) => {
+  return <div
       className={cn("animate-pulse rounded-md bg-primary/10", className)}
       aria-busy="true"
       role="progressbar"
       aria-valuemin={0}
       aria-valuemax={100}
       aria-valuetext="Please wait..."
       {...props}
+      ref={ref}
     />
-  )
-}
-
+})
 export { Skeleton }