2 files changed, 35 insertions, 15 deletions
diff --git a/extension/react-app/src/components/CodeBlock.tsx b/extension/react-app/src/components/CodeBlock.tsx
index 0aae0bbb..98760fbd 100644
--- a/extension/react-app/src/components/CodeBlock.tsx
+++ b/extension/react-app/src/components/CodeBlock.tsx
@@ -5,28 +5,27 @@ import { defaultBorderRadius, secondaryDark, vscBackground } from ".";
 
 import { Clipboard, CheckCircle } from "@styled-icons/heroicons-outline";
 
+import StyledCode from "./StyledCode";
+
 const StyledPre = styled.pre`
   overflow-y: scroll;
   word-wrap: normal;
   border: 0.5px solid gray;
   border-radius: ${defaultBorderRadius};
   background-color: ${secondaryDark};
-  padding: 8px;
-  padding-top: 14px;
-  padding-bottom: 16px;
+  position: relative;
 `;
 
-const StyledCode = styled.code``;
+const CopyButtonDiv = styled.div`
+  position: absolute;
+  top: 4px;
+  right: 4px;
+`;
 
 const StyledCopyButton = styled.button<{ visible: boolean }>`
-  /* position: relative; */
-  float: right;
   border: none;
-  background-color: ${secondaryDark};
+  background-color: transparent;
   cursor: pointer;
-  padding: 0;
-  /* margin: 4px; */
-  margin-top: -6px;
 
   visibility: ${(props) => (props.visible ? "visible" : "hidden")};
 `;
@@ -77,11 +76,13 @@ function CodeBlock(props: { language?: string; children: string }) {
         setHovered(false);
       }}
     >
-      <CopyButton
-        visible={hovered}
-        textToCopy={(props.children as any).props.children[0]}
-      />
-      <StyledCode>{props.children}</StyledCode>
+      <CopyButtonDiv>
+        <CopyButton
+          visible={hovered}
+          textToCopy={(props.children as any).props.children[0]}
+        />
+      </CopyButtonDiv>
+      <StyledCode language={props.language}>{props.children}</StyledCode>
     </StyledPre>
   );
 }
diff --git a/extension/react-app/src/components/StyledCode.tsx b/extension/react-app/src/components/StyledCode.tsx
new file mode 100644
index 00000000..485c70b6
--- /dev/null
+++ b/extension/react-app/src/components/StyledCode.tsx
@@ -0,0 +1,19 @@
+import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
+import { vscDarkPlus as highlightStyle } from "react-syntax-highlighter/dist/esm/styles/prism";
+
+interface StyledCodeProps {
+  children: string;
+  language?: string;
+}
+
+const StyledCode = (props: StyledCodeProps) => (
+  <SyntaxHighlighter
+    customStyle={{ margin: "0" }}
+    style={highlightStyle}
+    language={props.language || "python"}
+  >
+    {(props.children as any).props.children}
+  </SyntaxHighlighter>
+);
+
+export default StyledCode;