feat: add support for icons and warning variant in Badge component (#17350)

jaaydenh · web-flow · commit 6a6e1ec50c5d · 2025-04-11T10:16:37.000-04:00
diff --git a/site/src/components/Badge/Badge.stories.tsx b/site/src/components/Badge/Badge.stories.tsx
@@ -1,4 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react";
+import { Settings, TriangleAlert } from "lucide-react";
 import { Badge } from "./Badge";
 
 const meta: Meta<typeof Badge> = {
@@ -13,3 +14,25 @@ export default meta;
 type Story = StoryObj<typeof Badge>;
 
 export const Default: Story = {};
+
+export const Warning: Story = {
+	args: {
+		variant: "warning",
+	},
+};
+
+export const SmallWithIcon: Story = {
+	args: {
+		variant: "default",
+		size: "sm",
+		children: <>{<Settings />} Preset</>,
+	},
+};
+
+export const MediumWithIcon: Story = {
+	args: {
+		variant: "warning",
+		size: "md",
+		children: <>{<TriangleAlert />} Immutable</>,
+	},
+};
diff --git a/site/src/components/Badge/Badge.tsx b/site/src/components/Badge/Badge.tsx
@@ -2,21 +2,26 @@
  * Copied from shadc/ui on 11/13/2024
  * @see {@link https://ui.shadcn.com/docs/components/badge}
  */
+import { Slot } from "@radix-ui/react-slot";
 import { type VariantProps, cva } from "class-variance-authority";
-import type { FC } from "react";
+import { forwardRef } from "react";
 import { cn } from "utils/cn";
 
 export const badgeVariants = cva(
-	"inline-flex items-center rounded-md border px-2 py-1 transition-colors focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2",
+	`inline-flex items-center rounded-md border px-2 py-1 transition-colors
+	focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2
+	[&_svg]:pointer-events-none [&_svg]:pr-0.5 [&_svg]:py-0.5 [&_svg]:mr-0.5`,
 	{
 		variants: {
 			variant: {
 				default:
 					"border-transparent bg-surface-secondary text-content-secondary shadow",
+				warning:
+					"border-transparent bg-surface-orange text-content-warning shadow",
 			},
 			size: {
-				sm: "text-2xs font-regular",
-				md: "text-xs font-medium",
+				sm: "text-2xs font-regular h-5.5 [&_svg]:size-icon-xs",
+				md: "text-xs font-medium [&_svg]:size-icon-sm",
 			},
 		},
 		defaultVariants: {
@@ -28,18 +33,20 @@ export const badgeVariants = cva(
 
 export interface BadgeProps
 	extends React.HTMLAttributes<HTMLDivElement>,
-		VariantProps<typeof badgeVariants> {}
+		VariantProps<typeof badgeVariants> {
+	asChild?: boolean;
+}
 
-export const Badge: FC<BadgeProps> = ({
-	className,
-	variant,
-	size,
-	...props
-}) => {
-	return (
-		<div
-			className={cn(badgeVariants({ variant, size }), className)}
-			{...props}
-		/>
-	);
-};
+export const Badge = forwardRef<HTMLDivElement, BadgeProps>(
+	({ className, variant, size, asChild = false, ...props }, ref) => {
+		const Comp = asChild ? Slot : "div";
+
+		return (
+			<Comp
+				{...props}
+				ref={ref}
+				className={cn(badgeVariants({ variant, size }), className)}
+			/>
+		);
+	},
+);
