From b29aedf1c8537f03976b43da3a608bb8bbd64175 Mon Sep 17 00:00:00 2001
From: Jamie Kyle <jamie@signal.org>
Date: Wed, 12 Nov 2025 17:03:30 -0800
Subject: [PATCH] Add new axo checkbox square variant

---
 ts/axo/AxoCheckbox.dom.stories.tsx | 30 ++++++++++++++++++---
 ts/axo/AxoCheckbox.dom.tsx         | 43 ++++++++++++++++++++++++------
 ts/axo/AxoDialog.dom.stories.tsx   |  1 +
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/ts/axo/AxoCheckbox.dom.stories.tsx b/ts/axo/AxoCheckbox.dom.stories.tsx
index ec9fb4939e..fa58e7bd4a 100644
--- a/ts/axo/AxoCheckbox.dom.stories.tsx
+++ b/ts/axo/AxoCheckbox.dom.stories.tsx
@@ -11,6 +11,7 @@ export default {
 
 function Template(props: {
   label: string;
+  variant: AxoCheckbox.Variant;
   defaultChecked: boolean;
   disabled?: boolean;
 }): JSX.Element {
@@ -18,6 +19,7 @@ function Template(props: {
   return (
     <label className={tw('my-2 flex items-center gap-2')}>
       <AxoCheckbox.Root
+        variant={props.variant}
         checked={checked}
         onCheckedChange={setChecked}
         disabled={props.disabled}
@@ -31,10 +33,30 @@ export function Basic(): JSX.Element {
   return (
     <>
       <h1 className={tw('type-title-large')}>AxoCheckbox</h1>
-      <Template label="Unchecked" defaultChecked={false} />
-      <Template label="Checked" defaultChecked />
-      <Template label="Unchecked+Disabled" defaultChecked={false} disabled />
-      <Template label="Checked+Disabled" defaultChecked disabled />
+      {AxoCheckbox._getAllCheckboxVariants().map(variant => {
+        return (
+          <section>
+            <Template
+              variant={variant}
+              label="Unchecked"
+              defaultChecked={false}
+            />
+            <Template variant={variant} label="Checked" defaultChecked />
+            <Template
+              variant={variant}
+              label="Unchecked+Disabled"
+              defaultChecked={false}
+              disabled
+            />
+            <Template
+              variant={variant}
+              label="Checked+Disabled"
+              defaultChecked
+              disabled
+            />
+          </section>
+        );
+      })}
     </>
   );
 }
diff --git a/ts/axo/AxoCheckbox.dom.tsx b/ts/axo/AxoCheckbox.dom.tsx
index 6ccdb974df..9367c33aec 100644
--- a/ts/axo/AxoCheckbox.dom.tsx
+++ b/ts/axo/AxoCheckbox.dom.tsx
@@ -3,13 +3,37 @@
 import React, { memo } from 'react';
 import { Checkbox } from 'radix-ui';
 import { AxoSymbol } from './AxoSymbol.dom.js';
+import type { TailwindStyles } from './tw.dom.js';
 import { tw } from './tw.dom.js';
 
 const Namespace = 'AxoCheckbox';
 
 export namespace AxoCheckbox {
+  type VariantConfig = {
+    rootStyles: TailwindStyles;
+    iconSize: AxoSymbol.IconSize;
+  };
+
+  const Variants: Record<Variant, VariantConfig> = {
+    round: {
+      rootStyles: tw('size-5 rounded-full'),
+      iconSize: 14,
+    },
+    square: {
+      rootStyles: tw('size-4 rounded-sm'),
+      iconSize: 12,
+    },
+  };
+
+  export function _getAllCheckboxVariants(): ReadonlyArray<Variant> {
+    return Object.keys(Variants) as Array<Variant>;
+  }
+
+  export type Variant = 'round' | 'square';
+
   export type RootProps = Readonly<{
     id?: string;
+    variant: Variant;
     checked: boolean;
     onCheckedChange: (nextChecked: boolean) => void;
     disabled?: boolean;
@@ -17,6 +41,7 @@ export namespace AxoCheckbox {
   }>;
 
   export const Root = memo((props: RootProps) => {
+    const variantConfig = Variants[props.variant];
     return (
       <Checkbox.Root
         id={props.id}
@@ -25,24 +50,26 @@ export namespace AxoCheckbox {
         disabled={props.disabled}
         required={props.required}
         className={tw(
-          'flex size-5 items-center justify-center rounded-full leading-none',
+          variantConfig.rootStyles,
+          'flex items-center justify-center',
           'border border-border-primary inset-shadow-on-color',
           'data-[state=unchecked]:bg-fill-primary',
           'data-[state=unchecked]:pressed:bg-fill-primary-pressed',
           'data-[state=checked]:bg-color-fill-primary',
+          'data-[state=checked]:text-label-primary-on-color',
           'data-[state=checked]:pressed:bg-color-fill-primary-pressed',
           'data-[disabled]:border-border-secondary',
+          'data-[state=checked]:data-[disabled]:text-label-disabled-on-color',
           'outline-0 outline-border-focused focused:outline-[2.5px]',
           'overflow-hidden'
         )}
       >
-        <Checkbox.Indicator
-          className={tw(
-            'data-[state=checked]:text-label-primary-on-color',
-            'data-[state=checked]:data-[disabled]:text-label-disabled-on-color'
-          )}
-        >
-          <AxoSymbol.Icon symbol="check" size={14} label={null} />
+        <Checkbox.Indicator asChild>
+          <AxoSymbol.Icon
+            symbol="check"
+            size={variantConfig.iconSize}
+            label={null}
+          />
         </Checkbox.Indicator>
       </Checkbox.Root>
     );
diff --git a/ts/axo/AxoDialog.dom.stories.tsx b/ts/axo/AxoDialog.dom.stories.tsx
index 748a61fdac..df032f1c5d 100644
--- a/ts/axo/AxoDialog.dom.stories.tsx
+++ b/ts/axo/AxoDialog.dom.stories.tsx
@@ -252,6 +252,7 @@ function CheckboxField(props: { label: string }) {
     <div className={tw('flex gap-3 py-2.5')}>
       <AxoCheckbox.Root
         id={id}
+        variant="square"
         checked={checked}
         onCheckedChange={setChecked}
       />