From 455774ed5286f497f902e254e49471ca2802d74f Mon Sep 17 00:00:00 2001
From: David Chau <46395074+ddcc4@users.noreply.github.com>
Date: Mon, 10 Feb 2025 14:13:13 -0500
Subject: [PATCH] feat(protocol-designer): add declaration for `def run()` to
 generated Python (#17482)

# Overview

This adds the `def run(...)` declaration to the generated Python file.
AUTH-1092 We'll start filling in the `run()` function in future PRs.

## Test Plan and Hands on Testing

Updated unit tests. Looked at generated file with feature flag turned
on.

## Risk assessment

Low. This just affects the exported Python file hidden under a feature
flag.
---
 .../file-data/__tests__/createFile.test.ts    |  3 +++
 .../src/file-data/selectors/fileCreator.ts    |  8 +++++-
 .../src/file-data/selectors/pythonFile.ts     | 25 ++++++++++++++++++-
 step-generation/src/utils/pythonFormat.ts     |  6 +++++
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/protocol-designer/src/file-data/__tests__/createFile.test.ts b/protocol-designer/src/file-data/__tests__/createFile.test.ts
index 613f0e52f6a..527e1cccf20 100644
--- a/protocol-designer/src/file-data/__tests__/createFile.test.ts
+++ b/protocol-designer/src/file-data/__tests__/createFile.test.ts
@@ -119,6 +119,9 @@ requirements = {
     "robotType": "OT-2",
     "apiLevel": "2.23",
 }
+
+def run(protocol: protocol_api.ProtocolContext):
+    pass
 `.trimStart()
     )
   })
diff --git a/protocol-designer/src/file-data/selectors/fileCreator.ts b/protocol-designer/src/file-data/selectors/fileCreator.ts
index 9850b802dc7..a12547a3b70 100644
--- a/protocol-designer/src/file-data/selectors/fileCreator.ts
+++ b/protocol-designer/src/file-data/selectors/fileCreator.ts
@@ -26,7 +26,12 @@ import {
   getModulesLoadInfo,
   getPipettesLoadInfo,
 } from './utils'
-import { pythonImports, pythonMetadata, pythonRequirements } from './pythonFile'
+import {
+  pythonDefRun,
+  pythonImports,
+  pythonMetadata,
+  pythonRequirements,
+} from './pythonFile'
 
 import type { SecondOrderCommandAnnotation } from '@opentrons/shared-data/commandAnnotation/types'
 import type {
@@ -310,6 +315,7 @@ export const createPythonFile: Selector<string> = createSelector(
         pythonImports(),
         pythonMetadata(fileMetadata),
         pythonRequirements(robotType),
+        pythonDefRun(),
       ]
         .filter(section => section) // skip any blank sections
         .join('\n\n') + '\n'
diff --git a/protocol-designer/src/file-data/selectors/pythonFile.ts b/protocol-designer/src/file-data/selectors/pythonFile.ts
index dbf5e9c3bf9..e4d9e63eec0 100644
--- a/protocol-designer/src/file-data/selectors/pythonFile.ts
+++ b/protocol-designer/src/file-data/selectors/pythonFile.ts
@@ -1,7 +1,11 @@
 /** Generate sections of the Python file for fileCreator.ts */
 
 import { FLEX_ROBOT_TYPE, OT2_ROBOT_TYPE } from '@opentrons/shared-data'
-import { formatPyDict } from '@opentrons/step-generation'
+import {
+  formatPyDict,
+  indentPyLines,
+  PROTOCOL_CONTEXT_NAME,
+} from '@opentrons/step-generation'
 import type { FileMetadataFields } from '../types'
 import type { RobotType } from '@opentrons/shared-data'
 
@@ -46,3 +50,22 @@ export function pythonRequirements(robotType: RobotType): string {
   }
   return `requirements = ${formatPyDict(requirements)}`
 }
+
+export function pythonDefRun(): string {
+  const sections: string[] = [
+    // loadModules(),
+    // loadLabware(),
+    // loadInstruments(),
+    // defineLiquids(),
+    // loadLiquids(),
+    // stepCommands(),
+  ]
+  const functionBody =
+    sections
+      .filter(section => section) // skip empty sections
+      .join('\n\n') || 'pass'
+  return (
+    `def run(${PROTOCOL_CONTEXT_NAME}: protocol_api.ProtocolContext):\n` +
+    `${indentPyLines(functionBody)}`
+  )
+}
diff --git a/step-generation/src/utils/pythonFormat.ts b/step-generation/src/utils/pythonFormat.ts
index 6efc6557542..76dcb1809fc 100644
--- a/step-generation/src/utils/pythonFormat.ts
+++ b/step-generation/src/utils/pythonFormat.ts
@@ -1,5 +1,11 @@
 /** Utility functions for Python code generation. */
 
+/** The variable name for the ProtocolContext object in the run() function.
+ * Our docs call it `protocol`, which is slightly misleading since the object is not
+ * the protocol itself, but we'll try to stay consistent with the docs.
+ */
+export const PROTOCOL_CONTEXT_NAME = 'protocol'
+
 const INDENT = '    '
 
 /** Indent each of the lines in `text`. */