feat(frontend): use typed dict for resource types in python (#1869)

* feat(frontend): python typed dict resource types
+ filter pickers according to lang

* feat: use typed dict for AI gen

frontend/src/lib/components/EditorBar.svelte

@@ -31,10 +31,10 @@
 	import { SCRIPT_EDITOR_SHOW_EXPLORE_OTHER_SCRIPTS } from '$lib/consts'
 	import { createEventDispatcher } from 'svelte'
 	import { sendUserToast } from '$lib/toast'
-	import { getScriptByPath } from '$lib/scripts'
+	import { getScriptByPath, scriptLangToEditorLang } from '$lib/scripts'
 	import Toggle from './Toggle.svelte'
 	import { Link, Users } from 'lucide-svelte'
-	import { capitalize } from '$lib/utils'
+	import { capitalize, toCamel } from '$lib/utils'
 	import type { Schema, SchemaProperty, SupportedLanguage } from '$lib/common'
 	import ScriptVersionHistory from './ScriptVersionHistory.svelte'
 	import { ScriptGen } from './codeGen'
@@ -65,6 +65,15 @@
 	let resourceTypePicker: ItemPicker
 	let variableEditor: VariableEditor
 	let resourceEditor: ResourceEditor
+	let showContextVarPicker = false
+	let showVarPicker = false
+	let showResourcePicker = false
+	let showResourceTypePicker = false
+
+	$: showContextVarPicker = ['python3', 'bash', 'go', 'deno', 'bun'].includes(lang)
+	$: showVarPicker = ['python3', 'bash', 'go', 'deno', 'bun'].includes(lang)
+	$: showResourcePicker = ['python3', 'bash', 'go', 'deno', 'bun'].includes(lang)
+	$: showResourceTypePicker = scriptLangToEditorLang(lang) === 'typescript' || lang === 'python3'
 
 	let codeViewer: Drawer
 	let codeObj: { language: SupportedLanguage; content: string } | undefined = undefined
@@ -137,6 +146,35 @@
 		return rec(schema.properties, true)
 	}
 
+	function pythonCompile(schema: Schema) {
+		let res = ''
+		const entries = Object.entries(schema.properties)
+		if (entries.length === 0) {
+			return 'dict'
+		}
+		let i = 0
+		for (let [name, prop] of entries) {
+			let typ = 'dict'
+			if (prop.type === 'array') {
+				typ = 'list'
+			} else if (prop.type === 'string') {
+				typ = 'str'
+			} else if (prop.type === 'number') {
+				typ = 'float'
+			} else if (prop.type === 'integer') {
+				typ = 'int'
+			} else if (prop.type === 'boolean') {
+				typ = 'bool'
+			}
+			res += `${name}: ${typ}`
+			i++
+			if (i < entries.length) {
+				res += '\n'
+			}
+		}
+		return res
+	}
+
 	let historyBrowserDrawerOpen = false
 </script>
 
@@ -171,13 +209,14 @@
 		{/if}
 	</DrawerContent>
 </Drawer>
-
 <ItemPicker
 	bind:this={contextualVariablePicker}
 	pickCallback={(path, name) => {
 		if (!editor) return
 		if (lang == 'deno') {
 			editor.insertAtCursor(`Deno.env.get('${name}')`)
+		} else if (lang === 'bun') {
+			editor.insertAtCursor(`Bun.env["${name}"]`)
 		} else if (lang == 'python3') {
 			if (!editor.getCode().includes('import os')) {
 				editor.insertAtBeginning('import os\n')
@@ -212,6 +251,18 @@
 				)
 			}
 			editor.insertAtCursor(`(await wmill.getVariable('${path}'))`)
+		} else if (lang === 'bun') {
+			const code = editor.getCode()
+			if (!code.includes('import { setClient, getVariable } from "windmill-client@0.3.15')) {
+				editor.insertAtBeginning(
+					`import { setClient, getVariable } from "windmill-client@0.3.15"\n`
+				)
+			}
+			if (!code.includes('setClient()')) {
+				editor.insertAtCursor(`setClient()\n(await getVariable('${path}'))`)
+			} else {
+				editor.insertAtCursor(`(await getVariable('${path}'))`)
+			}
 		} else if (lang == 'python3') {
 			if (!editor.getCode().includes('import wmill')) {
 				editor.insertAtBeginning('import wmill\n')
@@ -262,6 +313,18 @@
 				)
 			}
 			editor.insertAtCursor(`(await wmill.getResource('${path}'))`)
+		} else if (lang === 'bun') {
+			const code = editor.getCode()
+			if (!code.includes('import { setClient, getResource } from "windmill-client@0.3.15')) {
+				editor.insertAtBeginning(
+					`import { setClient, getResource } from "windmill-client@0.3.15"\n`
+				)
+			}
+			if (!code.includes('setClient()')) {
+				editor.insertAtCursor(`setClient()\n(await getResource('${path}'))`)
+			} else {
+				editor.insertAtCursor(`(await getResource('${path}'))`)
+			}
 		} else if (lang == 'python3') {
 			if (!editor.getCode().includes('import wmill')) {
 				editor.insertAtBeginning('import wmill\n')
@@ -301,24 +364,29 @@
 	</div>
 </ItemPicker>
 
-{#if lang == 'deno'}
+{#if showResourceTypePicker}
 	<ItemPicker
 		bind:this={resourceTypePicker}
 		pickCallback={async (_, name) => {
 			if (!editor) return
-			const toCamel = (s) => {
-				return s.replace(/([-_][a-z])/gi, ($1) => {
-					return $1.toUpperCase().replace('-', '').replace('_', '')
-				})
-			}
 			const resourceType = await ResourceService.getResourceType({
 				workspace: $workspaceStore ?? 'NO_W',
 				path: name
 			})
 
-			const tsSchema = compile(resourceType.schema)
-			console.log(tsSchema)
-			editor.insertAtCursor(`type ${toCamel(capitalize(name))} = ${tsSchema}\n`)
+			if (lang == 'python3') {
+				const pySchema = pythonCompile(resourceType.schema)
+
+				editor.insertAtCursor(`class ${name}(TypedDict):\n${pySchema}\n`)
+				const code = editor.getCode()
+				if (!code.includes('from typing import TypedDict')) {
+					editor.insertAtBeginning('from typing import TypedDict\n')
+				}
+			} else {
+				const tsSchema = compile(resourceType.schema)
+				console.log(tsSchema)
+				editor.insertAtCursor(`type ${toCamel(capitalize(name))} = ${tsSchema}\n`)
+			}
 			sendUserToast(`${name} inserted at cursor`)
 		}}
 		tooltip="Resources Types are the schemas associated with a Resource. They define the structure of the data that is returned from a Resource."
@@ -339,46 +407,51 @@
 			class="rounded-full w-2 h-2 mx-2 {validCode ? 'bg-green-300' : 'bg-red-300'}"
 		/>
 		<div class="flex items-center">
-			<Button
-				title="Add context variable"
-				color="light"
-				btnClasses="!font-medium text-gray-600"
-				on:click={contextualVariablePicker.openDrawer}
-				size="xs"
-				spacingSize="md"
-				startIcon={{ icon: faDollarSign }}
-				{iconOnly}
-			>
-				+Context Var
-			</Button>
+			{#if showContextVarPicker}
+				<Button
+					title="Add context variable"
+					color="light"
+					btnClasses="!font-medium text-gray-600"
+					on:click={contextualVariablePicker.openDrawer}
+					size="xs"
+					spacingSize="md"
+					startIcon={{ icon: faDollarSign }}
+					{iconOnly}
+				>
+					+Context Var
+				</Button>
+			{/if}
+			{#if showVarPicker}
+				<Button
+					title="Add variable"
+					color="light"
+					btnClasses="!font-medium text-gray-600"
+					on:click={variablePicker.openDrawer}
+					size="xs"
+					spacingSize="md"
+					startIcon={{ icon: faDollarSign }}
+					{iconOnly}
+				>
+					+Variable
+				</Button>
+			{/if}
 
-			<Button
-				title="Add variable"
-				color="light"
-				btnClasses="!font-medium text-gray-600"
-				on:click={variablePicker.openDrawer}
-				size="xs"
-				spacingSize="md"
-				startIcon={{ icon: faDollarSign }}
-				{iconOnly}
-			>
-				+Variable
-			</Button>
+			{#if showResourcePicker}
+				<Button
+					title="Add resource"
+					btnClasses="!font-medium text-gray-600"
+					size="xs"
+					spacingSize="md"
+					color="light"
+					on:click={resourcePicker.openDrawer}
+					{iconOnly}
+					startIcon={{ icon: faCube }}
+				>
+					+Resource
+				</Button>
+			{/if}
 
-			<Button
-				title="Add resource"
-				btnClasses="!font-medium text-gray-600"
-				size="xs"
-				spacingSize="md"
-				color="light"
-				on:click={resourcePicker.openDrawer}
-				{iconOnly}
-				startIcon={{ icon: faCube }}
-			>
-				+Resource
-			</Button>
-
-			{#if lang == 'deno'}
+			{#if showResourceTypePicker}
 				<Button
 					title="Add resource"
 					btnClasses="!font-medium text-gray-600"

frontend/src/lib/components/codeGen/lib.ts

@@ -23,10 +23,11 @@ const COMMENT_TYPES = {
 const PROMPTS = {
 	[Script.language
 		.PYTHON3]: `Write a function in python called "main". The function should {description}. Specify the parameter types. Do not call the main function.
-You have access to the following resource types, if you need them, you have to define a TypedDict with the name specified (DO NOT CAPITALIZE) and add them as parameters: {resourceTypes}`,
+You have access to the following resource types, if you need them, you have to define the TypedDict exactly as specified (class name has to be IN LOWERCASE) and add them as parameters: {resourceTypes}
+If the TypedDict name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
 	[Script.language
 		.DENO]: `Write a function in typescript called "main". The function should {description}. Specify the parameter types. You are in a Deno environment. You can import deno libraries or you can also import npm libraries like that: "import ... from "npm:{package}";". Export the "main" function like this: "export function main(...)". Do not call the main function.
-You have access to the following resource types, if you need them, you have to define the type with the name specified and add them as parameters: {resourceTypes}
+You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
 If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
 	[Script.language.GO]:
 		'Write a function in go called "main". The function should {description}. Import the packages you need. The return type of the function has to be ({return_type}, error). The file package has to be "inner".',
@@ -38,11 +39,11 @@ If the type name conflicts with the imported object, rename the imported object
 		'Write SQL code for MySQL that should {description}. Arguments can be obtained directly in the statement with ?. Name the parameters by adding comments before the command like that: -- ? name1 ({type}) (one per row)',
 	[Script.language
 		.NATIVETS]: `Write a function in typescript called "main". The function should {description}. Specify the parameter types. You should use fetch and are not allowed to import any libraries. Export the "main" function like this: "export function main(...)". Do not call the main function.
-You have access to the following resource types, if you need them, you have to define the type with the name specified and add them as parameters: {resourceTypes}
+You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
 If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`,
 	[Script.language
 		.BUN]: `Write a function in typescript called "main".  The function should {description}. Specify the parameter types. You are in a Node.js environment. You can import npm libraries. Export the "main" function like this: "export function main(...)". Do not call the main function.
-You have access to the following resource types, if you need them, you have to define the type with the name specified and add them as parameters: {resourceTypes}
+You have access to the following resource types, if you need them, you have to define the type exactly as specified and add them as parameters: {resourceTypes}
 If the type name conflicts with the imported object, rename the imported object NOT THE TYPE.`
 }
 
@@ -129,11 +130,12 @@ export async function generateScript(scriptOptions: ScriptGenerationOptions): Pr
 		)
 	) {
 		const resourceTypes = await ResourceService.listResourceType({ workspace })
-		const resourceTypesText = formatResourceTypes(resourceTypes, true)
+		const resourceTypesText = formatResourceTypes(resourceTypes, 'typescript')
+
 		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
 	} else if (scriptOptions.language == Script.language.PYTHON3) {
 		const resourceTypes = await ResourceService.listResourceType({ workspace })
-		const resourceTypesText = formatResourceTypes(resourceTypes, false)
+		const resourceTypesText = formatResourceTypes(resourceTypes, 'python3')
 		prompt = prompt.replace('{resourceTypes}', resourceTypesText)
 	}
 

frontend/src/lib/components/codeGen/utils.ts

@@ -2,11 +2,11 @@ import type { Schema, SchemaProperty } from '../../common'
 
 import type { ResourceType } from '../../gen'
 
-import { capitalize } from '$lib/utils'
+import { capitalize, toCamel } from '$lib/utils'
 
 function compile(schema: Schema) {
 	function rec(x: { [name: string]: SchemaProperty }, root = false) {
-		let res = '{ '
+		let res = '{\n'
 		const entries = Object.entries(x)
 		if (entries.length == 0) {
 			return 'any'
@@ -14,39 +14,66 @@ function compile(schema: Schema) {
 		let i = 0
 		for (let [name, prop] of entries) {
 			if (prop.type == 'object') {
-				res += `${name}: ${rec(prop.properties ?? {})}`
+				res += `  ${name}: ${rec(prop.properties ?? {})}`
 			} else if (prop.type == 'array') {
-				res += `${name}: ${prop?.items?.type ?? 'any'}[]`
+				res += `  ${name}: ${prop?.items?.type ?? 'any'}[]`
 			} else {
 				let typ = prop?.type ?? 'any'
 				if (typ == 'integer') {
 					typ = 'number'
 				}
-				res += `${name}: ${typ}`
+				res += `  ${name}: ${typ}`
 			}
 			i++
 			if (i < entries.length) {
-				res += ', '
+				res += ',\n'
 			}
 		}
-		res += ' }'
+		res += '\n}'
 		return res
 	}
 	return rec(schema.properties, true)
 }
 
-function toCamel(s) {
-	return s.replace(/([-_][a-z])/gi, ($1) => {
-		return $1.toUpperCase().replace('-', '').replace('_', '')
-	})
+export function pythonCompile(schema: Schema) {
+	let res = ''
+	const entries = Object.entries(schema.properties)
+	if (entries.length === 0) {
+		return 'dict'
+	}
+	let i = 0
+	for (let [name, prop] of entries) {
+		let typ = 'dict'
+		if (prop.type === 'array') {
+			typ = 'list'
+		} else if (prop.type === 'string') {
+			typ = 'str'
+		} else if (prop.type === 'number') {
+			typ = 'float'
+		} else if (prop.type === 'integer') {
+			typ = 'int'
+		} else if (prop.type === 'boolean') {
+			typ = 'bool'
+		}
+		res += `    ${name}: ${typ}`
+		i++
+		if (i < entries.length) {
+			res += '\n'
+		}
+	}
+	return res
 }
 
-export function formatResourceTypes(resourceTypes: ResourceType[], camel = false) {
-	const result = resourceTypes.map((resourceType) => {
-		return `${camel ? toCamel(capitalize(resourceType.name)) : resourceType.name} ${compile(
-			resourceType.schema
-		)}`
-	})
-
-	return result.join(', ')
+export function formatResourceTypes(resourceTypes: ResourceType[], lang: 'python3' | 'typescript') {
+	if (lang === 'python3') {
+		const result = resourceTypes.map((resourceType) => {
+			return `class ${resourceType.name}(TypedDict):\n${pythonCompile(resourceType.schema)}`
+		})
+		return '\n' + result.join('\n\n')
+	} else {
+		const result = resourceTypes.map((resourceType) => {
+			return `type ${toCamel(capitalize(resourceType.name))} = ${compile(resourceType.schema)}`
+		})
+		return '\n' + result.join('\n\n')
+	}
 }

frontend/src/lib/utils.ts

@@ -579,3 +579,9 @@ export function extractCustomProperties(styleStr: string): string {
 
 	return customStyleStr
 }
+
+export function toCamel(s: string) {
+	return s.replace(/([-_][a-z])/gi, ($1) => {
+		return $1.toUpperCase().replace('-', '').replace('_', '')
+	})
+}