diff --git a/src/mdx/code.js b/src/mdx/code.js
index c68a1b2f0e6..1cbf73d9f3f 100644
--- a/src/mdx/code.js
+++ b/src/mdx/code.js
@@ -43,6 +43,11 @@ export const InlineCode = styled.code`
   background-color: ${themeGet('colors.neutral.muted')};
   border-radius: ${themeGet('radii.2')};
 `
+const colorMap = {
+  'token comment': '#747458',
+  'token function': '#cf3846',
+  'token parameter variable': '#277d7b',
+}
 
 const MonoText = props => <Text sx={{fontFamily: 'mono', fontSize: 1}} {...props} />
 
@@ -113,18 +118,14 @@ function Code({className = '', prompt, children}) {
         <CodeBlock className={highlightClassName} style={style} code={code}>
           {tokens.map((line, i) => (
             <Box key={i} {...getLineProps({line, key: i})}>
-              {line.map((token, key) => (
-                <MonoText
-                  key={key}
-                  {...{
-                    ...getTokenProps({token, key}),
-                    style:
-                      getTokenProps({token, key}).className === 'token comment'
-                        ? {...getTokenProps({token, key}).style, color: '#747458'}
-                        : getTokenProps({token, key}).style,
-                  }}
-                />
-              ))}
+              {line.map((token, key) => {
+                const tokenProps = getTokenProps({token, key})
+                const tokenStyle = colorMap[tokenProps.className]
+                  ? {...tokenProps.style, color: colorMap[tokenProps.className]}
+                  : tokenProps.style
+
+                return <MonoText key={key} {...tokenProps} style={tokenStyle} />
+              })}
             </Box>
           ))}
         </CodeBlock>