feat: copy button in dag nodes (#696)

Co-authored-by: Cahllagerfeld <[email protected]>

Author: Aagya-Sharma

src/components/dag-visualizer/ArtifactNode.tsx

@@ -1,9 +1,10 @@
+import Copy from "@/assets/icons/copy.svg?react";
 import { ArtifactVersion } from "@/types/artifact-versions";
-
 import { NodeProps, useStore } from "reactflow";
 import { ArtifactIcon } from "../ArtifactIcon";
 import { ArtifactSheet } from "../artifacts/artifact-node-sheet";
 import { BaseNode } from "./BaseNode";
+import { CopyNodeButton } from "./NodeCopyButton";
 
 export function ArtifactNode({ data, selected }: NodeProps<ArtifactVersion & { name: string }>) {
 	const { unselectNodesAndEdges } = useStore((state) => ({
@@ -42,11 +43,23 @@ export function ArtifactNode({ data, selected }: NodeProps<ArtifactVersion & { n
 						<p className="truncate text-text-sm font-semibold text-theme-text-brand group-data-[selected=true]:text-theme-text-negative">
 							{data.name}
 						</p>
+
 						<p className="truncate text-text-xs text-theme-text-secondary group-data-[selected=true]:text-white/70">
 							{/* As artifact_type doesn't correspond to the last part of the string */}
 							{getTypeFromArtifact(data.body?.data_type.attribute || "")}
 						</p>
 					</div>
+					<CopyNodeButton
+						className="h-4 w-4 shrink-0 rounded-sm hover:bg-primary-100 active:bg-primary-200"
+						code={`from zenml.client import Client
+							
+artifact = Client().get_artifact_version('${data.id}')
+loaded_artifact = artifact.load()`}
+						type="artifact"
+					>
+						<Copy className="h-3 w-3 fill-primary-400" />
+						<div className="sr-only">Copy code to load artifact</div>
+					</CopyNodeButton>
 				</button>
 			</ArtifactSheet>
 		</BaseNode>

src/components/dag-visualizer/NodeCopyButton.tsx

@@ -0,0 +1,53 @@
+import { HTMLAttributes, useState } from "react";
+import {
+	Tooltip,
+	TooltipContent,
+	TooltipProvider,
+	TooltipTrigger
+} from "@zenml-io/react-component-library/components/client";
+import { cn } from "@zenml-io/react-component-library/utilities";
+
+export function CopyNodeButton({
+	className,
+	code,
+	type,
+	...rest
+}: HTMLAttributes<HTMLDivElement> & { code: string; type: "artifact" | "step" }) {
+	const [copied, setCopied] = useState(false);
+	const [tooltipOpen, setTooltipOpen] = useState(false);
+
+	function copyToClipboard(code: string) {
+		if (navigator.clipboard) {
+			navigator.clipboard.writeText(code);
+			setCopied(true);
+			setTooltipOpen(true);
+
+			setTimeout(() => {
+				setTooltipOpen(false);
+				setCopied(false);
+			}, 2000);
+		}
+	}
+	return (
+		<TooltipProvider>
+			<Tooltip open={tooltipOpen} onOpenChange={setTooltipOpen}>
+				<TooltipTrigger asChild>
+					<div
+						{...rest}
+						className={cn(
+							"hidden cursor-pointer items-center justify-center transition-all group-hover:flex",
+							className
+						)}
+						onClick={(e) => {
+							e.stopPropagation(); // Prevent the click event from bubbling up
+							copyToClipboard(code);
+						}}
+					></div>
+				</TooltipTrigger>
+				<TooltipContent className="w-full max-w-md whitespace-normal transition-opacity">
+					{copied ? "Copied!" : `Copy code to load ${type}`}
+				</TooltipContent>
+			</Tooltip>
+		</TooltipProvider>
+	);
+}

src/components/dag-visualizer/StepNode.tsx

@@ -1,10 +1,12 @@
+import Copy from "@/assets/icons/copy.svg?react";
 import { calculateTimeDifference } from "@/lib/dates";
 import { Step } from "@/types/steps";
 import { clsx } from "clsx";
 import { NodeProps, ReactFlowState, useStore } from "reactflow";
 import { ExecutionStatusIcon, getExecutionStatusBackgroundColor } from "../ExecutionStatus";
 import { StepSheet } from "../steps/step-sheet";
 import { BaseNode } from "./BaseNode";
+import { CopyNodeButton } from "./NodeCopyButton";
 
 const selector = (state: ReactFlowState) => ({
 	unselectAll: state.unselectNodesAndEdges
@@ -30,7 +32,7 @@ export function StepNode({ data, selected }: NodeProps<Step>) {
 					data-failed={isFailed}
 					data-selected={!!selected}
 					className={clsx(
-						"max-h-[80px] max-w-[300px] overflow-hidden rounded-md border bg-theme-surface-primary transition-all duration-200  hover:shadow-md  data-[selected=true]:shadow-md",
+						"group max-h-[80px] max-w-[300px] overflow-hidden rounded-md border bg-theme-surface-primary transition-all duration-200  hover:shadow-md  data-[selected=true]:shadow-md",
 						{
 							"border-theme-border-moderate hover:border-neutral-400 data-[selected=true]:border-theme-border-bold":
 								!isFailed,
@@ -45,6 +47,16 @@ export function StepNode({ data, selected }: NodeProps<Step>) {
 							<ExecutionStatusIcon status={data.body?.status} className="h-4 w-4" />
 						</div>
 						<p className="truncate font-semibold">{data.name}</p>
+						<CopyNodeButton
+							className="h-4 w-4 shrink-0 rounded-sm hover:bg-theme-surface-secondary active:bg-neutral-300"
+							code={`from zenml.client import Client
+step = Client().get_run_step(${data.id})
+config = step.config`}
+							type="step"
+						>
+							<Copy className="h-3 w-3 fill-theme-text-tertiary" />
+							<div className="sr-only">Copy code to load step</div>
+						</CopyNodeButton>
 					</div>
 					<div
 						data-failed={isFailed}