feat(protocol-designer): load module in python file (#17483)

addresses a bit of AUTH-1092

protocol-designer/src/file-data/__tests__/createFile.test.ts

@@ -70,6 +70,12 @@ describe('createFile selector', () => {
   afterEach(() => {
     vi.restoreAllMocks()
   })
+  const entities = {
+    moduleEntities: v7Fixture.moduleEntities,
+    labwareEntities,
+    pipetteEntities,
+    liquidEntities: ingredients,
+  }
   it('should return a schema-valid JSON V8 protocol', () => {
     // @ts-expect-error(sa, 2021-6-15): resultFunc not part of Selector type
     const result = createFile.resultFunc(
@@ -78,16 +84,13 @@ describe('createFile selector', () => {
       v7Fixture.robotStateTimeline,
       OT2_ROBOT_TYPE,
       dismissedWarnings,
-      ingredients,
       ingredLocations,
       v7Fixture.savedStepForms,
       v7Fixture.orderedStepIds,
-      labwareEntities,
-      v7Fixture.moduleEntities,
-      pipetteEntities,
       labwareNicknamesById,
       labwareDefsByURI,
-      {}
+      {},
+      entities
     )
     expectResultToMatchSchema(result)
 
@@ -99,7 +102,12 @@ describe('createFile selector', () => {
 
   it('should return a valid Python protocol file', () => {
     // @ts-expect-error(sa, 2021-6-15): resultFunc not part of Selector type
-    const result = createPythonFile.resultFunc(fileMetadata, OT2_ROBOT_TYPE, {})
+    const result = createPythonFile.resultFunc(
+      fileMetadata,
+      OT2_ROBOT_TYPE,
+      entities,
+      v7Fixture.initialRobotState
+    )
     // This is just a quick smoke test to make sure createPythonFile() produces
     // something that looks like a Python file. The individual sections of the
     // generated Python will be tested in separate unit tests.

protocol-designer/src/file-data/__tests__/pythonFile.test.ts

@@ -1,6 +1,19 @@
 import { describe, it, expect } from 'vitest'
-import { FLEX_ROBOT_TYPE, OT2_ROBOT_TYPE } from '@opentrons/shared-data'
-import { pythonMetadata, pythonRequirements } from '../selectors/pythonFile'
+import {
+  FLEX_ROBOT_TYPE,
+  HEATERSHAKER_MODULE_TYPE,
+  HEATERSHAKER_MODULE_V1,
+  MAGNETIC_BLOCK_TYPE,
+  MAGNETIC_BLOCK_V1,
+  OT2_ROBOT_TYPE,
+} from '@opentrons/shared-data'
+import {
+  getLoadModules,
+  pythonMetadata,
+  pythonRequirements,
+} from '../selectors/pythonFile'
+import type { TimelineFrame } from '@opentrons/step-generation'
+import type { ModuleEntities } from '../../step-forms'
 
 describe('pythonMetadata', () => {
   it('should generate metadata section', () => {
@@ -50,3 +63,43 @@ requirements = {
     )
   })
 })
+
+describe('getLoadModules', () => {
+  it('should generate loadModules', () => {
+    const moduleId = '1'
+    const moduleId2 = '2'
+    const moduleId3 = '3'
+    const mockModuleEntities: ModuleEntities = {
+      [moduleId]: {
+        id: moduleId,
+        model: MAGNETIC_BLOCK_V1,
+        type: MAGNETIC_BLOCK_TYPE,
+        pythonName: 'magnetic_block_1',
+      },
+      [moduleId2]: {
+        id: moduleId2,
+        model: HEATERSHAKER_MODULE_V1,
+        type: HEATERSHAKER_MODULE_TYPE,
+        pythonName: 'heater_shaker_1',
+      },
+      [moduleId3]: {
+        id: moduleId3,
+        model: MAGNETIC_BLOCK_V1,
+        type: MAGNETIC_BLOCK_TYPE,
+        pythonName: 'magnetic_block_2',
+      },
+    }
+    const modules: TimelineFrame['modules'] = {
+      [moduleId]: { slot: 'B1', moduleState: {} as any },
+      [moduleId2]: { slot: 'A1', moduleState: {} as any },
+      [moduleId3]: { slot: 'A2', moduleState: {} as any },
+    }
+
+    expect(getLoadModules(mockModuleEntities, modules)).toBe(
+      `# Load Modules:
+magnetic_block_1 = protocol.load_module("magneticBlockV1", "B1")
+heater_shaker_1 = protocol.load_module("heaterShakerModuleV1", "A1")
+magnetic_block_2 = protocol.load_module("magneticBlockV1", "A2")`
+    )
+  })
+})

protocol-designer/src/file-data/selectors/fileCreator.ts

@@ -99,34 +99,34 @@ export const createFile: Selector<ProtocolFile> = createSelector(
   getRobotStateTimeline,
   getRobotType,
   dismissSelectors.getAllDismissedWarnings,
-  stepFormSelectors.getLiquidEntities,
   ingredSelectors.getLiquidsByLabwareId,
   stepFormSelectors.getSavedStepForms,
   stepFormSelectors.getOrderedStepIds,
-  stepFormSelectors.getLabwareEntities,
-  stepFormSelectors.getModuleEntities,
-  stepFormSelectors.getPipetteEntities,
   uiLabwareSelectors.getLabwareNicknamesById,
   labwareDefSelectors.getLabwareDefsByURI,
   getStepGroups,
+  stepFormSelectors.getInvariantContext,
   (
     fileMetadata,
     initialRobotState,
     robotStateTimeline,
     robotType,
     dismissedWarnings,
-    liquidEntities,
     ingredLocations,
     savedStepForms,
     orderedStepIds,
-    labwareEntities,
-    moduleEntities,
-    pipetteEntities,
     labwareNicknamesById,
     labwareDefsByURI,
-    stepGroups
+    stepGroups,
+    invariantContext
   ) => {
     const { author, description, created } = fileMetadata
+    const {
+      pipetteEntities,
+      labwareEntities,
+      liquidEntities,
+      moduleEntities,
+    } = invariantContext
 
     const loadCommands = getLoadCommands(
       initialRobotState,
@@ -308,14 +308,16 @@ export const createFile: Selector<ProtocolFile> = createSelector(
 export const createPythonFile: Selector<string> = createSelector(
   getFileMetadata,
   getRobotType,
-  (fileMetadata, robotType) => {
+  stepFormSelectors.getInvariantContext,
+  getInitialRobotState,
+  (fileMetadata, robotType, invariantContext, robotState) => {
     return (
       [
         // Here are the sections of the Python file:
         pythonImports(),
         pythonMetadata(fileMetadata),
         pythonRequirements(robotType),
-        pythonDefRun(),
+        pythonDefRun(invariantContext, robotState),
       ]
         .filter(section => section) // skip any blank sections
         .join('\n\n') + '\n'

protocol-designer/src/file-data/selectors/pythonFile.ts

@@ -3,11 +3,17 @@
 import { FLEX_ROBOT_TYPE, OT2_ROBOT_TYPE } from '@opentrons/shared-data'
 import {
   formatPyDict,
+  formatPyStr,
   indentPyLines,
   PROTOCOL_CONTEXT_NAME,
 } from '@opentrons/step-generation'
-import type { FileMetadataFields } from '../types'
+import type {
+  InvariantContext,
+  ModuleEntities,
+  TimelineFrame,
+} from '@opentrons/step-generation'
 import type { RobotType } from '@opentrons/shared-data'
+import type { FileMetadataFields } from '../types'
 
 const PAPI_VERSION = '2.23' // latest version from api/src/opentrons/protocols/api_support/definitions.py
 
@@ -51,9 +57,36 @@ export function pythonRequirements(robotType: RobotType): string {
   return `requirements = ${formatPyDict(requirements)}`
 }
 
-export function pythonDefRun(): string {
+export function getLoadModules(
+  moduleEntities: ModuleEntities,
+  moduleRobotState: TimelineFrame['modules']
+): string {
+  const hasModules = Object.keys(moduleEntities).length > 0
+  const pythonModules = hasModules
+    ? Object.values(moduleEntities)
+        .map(module => {
+          // pythonIdentifier (module.model) from api/src/opentrons/protocol_api/validation.py#L373
+          return `${
+            module.pythonName
+          } = ${PROTOCOL_CONTEXT_NAME}.load_module(${formatPyStr(
+            module.model
+          )}, ${formatPyStr(moduleRobotState[module.id].slot)})`
+        })
+        .join('\n')
+    : ''
+  return hasModules ? `# Load Modules:\n${pythonModules}` : ''
+}
+
+export function pythonDefRun(
+  invariantContext: InvariantContext,
+  robotState: TimelineFrame
+): string {
+  const { moduleEntities } = invariantContext
+
+  const loadModules = getLoadModules(moduleEntities, robotState.modules)
+
   const sections: string[] = [
-    // loadModules(),
+    loadModules,
     // loadLabware(),
     // loadInstruments(),
     // defineLiquids(),