From 557533baed6625b3339dee528c283396b87e5470 Mon Sep 17 00:00:00 2001
From: zach <zach@dylibso.com>
Date: Thu, 26 Sep 2024 13:39:06 -0700
Subject: [PATCH] fix: add doc comments to exports/implementation functions,
 fix pyright type checking

---
 template/plugin/__init__.py.ejs    | 6 +++++-
 template/plugin/pdk_imports.py.ejs | 2 +-
 template/plugin/plugin.py.ejs      | 3 +++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/template/plugin/__init__.py.ejs b/template/plugin/__init__.py.ejs
index 2347b8f..faa7f22 100644
--- a/template/plugin/__init__.py.ejs
+++ b/template/plugin/__init__.py.ejs
@@ -16,13 +16,17 @@ from pdk_types import <%- Object.values(schema.schemas).map(schema => schema.nam
 # And it returns an output <%- toPythonType(imp.output) %> (<%- formatCommentLine(imp.output.description) %>)
 <% } -%>
 @extism.import_fn("extism:host/user", "<%- imp.name %>")
-def <%- camelToSnakeCase(imp.name) %>(<% if (imp.input) { -%>input: <%- toPythonType(imp.input) %><%} -%>) <% if (imp.output) { %> -> <%- toPythonType(imp.output) %><% } %>:
+def <%- camelToSnakeCase(imp.name) %>(<% if (imp.input) { -%>input: <%- toPythonType(imp.input) %><%} -%>) <% if (imp.output) { %> -> <%- toPythonType(imp.output) %><% } %>: # pyright: ignore [reportReturnType]
   pass
 <% }) %>
 
 # Exports
+# The implementations for these functions is in `plugin.py`
 
 <% schema.exports.forEach(ex => { -%>
+<% if (hasComment(ex)) { -%>
+# <%- formatCommentBlock(ex.description, "# ") %>
+<% } -%>
 @extism.plugin_fn
 def <%- ex.name %>():
   res = plugin.<%- camelToSnakeCase(ex.name) %>(<% if (ex.input) { %> extism.input(<%- toPythonType(ex.input) %>) <% } %>)
diff --git a/template/plugin/pdk_imports.py.ejs b/template/plugin/pdk_imports.py.ejs
index 96b340f..7cdde44 100644
--- a/template/plugin/pdk_imports.py.ejs
+++ b/template/plugin/pdk_imports.py.ejs
@@ -13,7 +13,7 @@ from pdk_types import <%- Object.values(schema.schemas).map(schema => schema.nam
 # And it returns an output <%- toPythonType(imp.output) %> (<%- formatCommentLine(imp.output.description) %>)
 <% } -%>
 @extism.import_fn("extism:host/user", "<%- imp.name %>")
-def <%- camelToSnakeCase(imp.name) %>(<% if (imp.input) { -%>input: <%- toPythonType(imp.input) %><%} -%>) <% if (imp.output) { %> -> <%- toPythonType(imp.output) %><% } %>:
+def <%- camelToSnakeCase(imp.name) %>(<% if (imp.input) { -%>input: <%- toPythonType(imp.input) %><%} -%>) <% if (imp.output) { %> -> <%- toPythonType(imp.output) %><% } %>: # pyright: ignore [reportReturnType]
   pass
 <% }) %>
 
diff --git a/template/plugin/plugin.py.ejs b/template/plugin/plugin.py.ejs
index 6002585..05ae433 100644
--- a/template/plugin/plugin.py.ejs
+++ b/template/plugin/plugin.py.ejs
@@ -4,6 +4,9 @@ from pdk_imports import <%- schema.imports.map(schema => camelToSnakeCase(schema
 from typing import List, Optional # noqa: F401
 
 <% schema.exports.forEach(ex => { -%>
+<% if (hasComment(ex)) { -%>
+# <%- formatCommentBlock(ex.description, "# ") %>
+<% } -%>
 def <%- camelToSnakeCase(ex.name) %>(<% if (ex.input) { %>input: <%- toPythonType(ex.input) %> <% } %>) <% if (ex.output) {%>-> <%- toPythonType(ex.output) %><%}%>:
 <% if (featureFlags['stub-with-code-samples'] && codeSamples(ex, 'python').length > 0) { -%>
 	<%- codeSamples(ex, 'python')[0].source %>