Enable TextInputWithTokens to accept leading and trailing visuals

.changeset/hot-toes-count.md

@@ -0,0 +1,5 @@
+---
+'@primer/react': patch
+---
+
+Makes it possible to render leading and trailing visuals in TextInputWithTokens just like we do in TextInput

docs/content/TextInputWithTokens.mdx

@@ -173,7 +173,7 @@ render(MaxHeightExample)
 ### With an error validation status
 
 ```javascript live noinline
-const BasicExample = () => {
+const ErrorExample = () => {
   const [tokens, setTokens] = React.useState([
     {text: 'zero', id: 0},
     {text: 'one', id: 1},
@@ -198,7 +198,45 @@ const BasicExample = () => {
   )
 }
 
-render(BasicExample)
+render(ErrorExample)
+```
+
+### With leading and trailing visuals
+
+```javascript live noinline
+const LeadingVisualExample = () => {
+  const [dates, setDates] = React.useState([
+    {text: '01 Jan', id: 0},
+    {text: '01 Feb', id: 1},
+    {text: '01 Mar', id: 2}
+  ])
+  const [tokens, setTokens] = React.useState([
+    {text: 'zero', id: 0},
+    {text: 'one', id: 1},
+    {text: 'two', id: 2}
+  ])
+  const onDateRemove = tokenId => {
+    setDates(dates.filter(token => token.id !== tokenId))
+  }
+  const onTokenRemove = tokenId => {
+    setTokens(tokens.filter(token => token.id !== tokenId))
+  }
+
+  return (
+    <>
+      <FormControl>
+        <FormControl.Label>Dates</FormControl.Label>
+        <TextInputWithTokens leadingVisual={CalendarIcon} tokens={dates} onTokenRemove={onDateRemove} />
+      </FormControl>
+      <FormControl>
+        <FormControl.Label>Tokens</FormControl.Label>
+        <TextInputWithTokens trailingVisual={CheckIcon} tokens={tokens} onTokenRemove={onTokenRemove} />
+      </FormControl>
+    </>
+  )
+}
+
+render(LeadingVisualExample)
 ```
 
 ## Props

src/TextInputWithTokens.tsx

@@ -66,6 +66,8 @@ const overflowCountFontSizeMap: Record<TokenSizeKeys, number> = {
 function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactComponent>(
   {
     icon: IconComponent,
+    leadingVisual: LeadingVisual,
+    trailingVisual: TrailingVisual,
     contrast,
     className,
     block,
@@ -248,7 +250,8 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
       className={className}
       contrast={contrast}
       disabled={disabled}
-      hasLeadingVisual={Boolean(IconComponent)}
+      hasLeadingVisual={Boolean(LeadingVisual)}
+      hasTrailingVisual={Boolean(TrailingVisual)}
       theme={theme}
       width={widthProp}
       minWidth={minWidthProp}
@@ -283,6 +286,12 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
         ...sxProp
       }}
     >
+      {IconComponent && !LeadingVisual && <IconComponent className="TextInput-icon" />}
+      {LeadingVisual && !IconComponent && (
+        <span className="TextInput-icon">
+          {typeof LeadingVisual === 'function' ? <LeadingVisual /> : LeadingVisual}
+        </span>
+      )}
       <Box
         ref={containerRef as RefObject<HTMLDivElement>}
         display="flex"
@@ -300,7 +309,6 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
           }
         }}
       >
-        {IconComponent && <IconComponent className="TextInput-icon" />}
         <Box
           sx={{
             order: 1,
@@ -344,6 +352,11 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
           </Text>
         ) : null}
       </Box>
+      {TrailingVisual && (
+        <span className="TextInput-icon">
+          {typeof TrailingVisual === 'function' ? <TrailingVisual /> : TrailingVisual}
+        </span>
+      )}
     </TextInputWrapper>
   )
 }

src/__tests__/TextInputWithTokens.test.tsx

@@ -6,6 +6,7 @@ import 'babel-polyfill'
 import {TokenSizeKeys, tokenSizes} from '../Token/TokenBase'
 import {IssueLabelToken} from '../Token'
 import TextInputWithTokens, {TextInputWithTokensProps} from '../TextInputWithTokens'
+import {MarkGithubIcon} from '@primer/octicons-react'
 expect.extend(toHaveNoViolations)
 
 const mockTokens = [
@@ -94,6 +95,20 @@ describe('TextInputWithTokens', () => {
     ).toMatchSnapshot()
   })
 
+  it('renders a leadingVisual and trailingVisual', () => {
+    const onRemoveMock = jest.fn()
+    expect(
+      render(
+        <TextInputWithTokens
+          leadingVisual={MarkGithubIcon}
+          trailingVisual={MarkGithubIcon}
+          tokens={mockTokens}
+          onTokenRemove={onRemoveMock}
+        />
+      )
+    ).toMatchSnapshot()
+  })
+
   it('focuses the previous token when keying ArrowLeft', () => {
     const onRemoveMock = jest.fn()
     const {getByLabelText, getByText} = HTMLRender(