From 3b4192007b71d5a99b2256283c32e92e9c4240de Mon Sep 17 00:00:00 2001
From: Lucas McDonald <lucasmcdonald3@gmail.com>
Date: Mon, 16 Sep 2024 13:52:31 -0700
Subject: [PATCH] feat(Python): Format docstrings

---
 .../python/codegen/DirectedPythonCodegen.java      | 12 ++++++++++++
 .../DafnyPythonLocalServiceStructureGenerator.java | 14 --------------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/codegen/smithy-dafny-codegen-modules/smithy-python/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/DirectedPythonCodegen.java b/codegen/smithy-dafny-codegen-modules/smithy-python/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/DirectedPythonCodegen.java
index 9592023fa..c415e0a6f 100644
--- a/codegen/smithy-dafny-codegen-modules/smithy-python/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/DirectedPythonCodegen.java
+++ b/codegen/smithy-dafny-codegen-modules/smithy-python/codegen/smithy-python-codegen/src/main/java/software/amazon/smithy/python/codegen/DirectedPythonCodegen.java
@@ -338,6 +338,7 @@ public void customizeAfterIntegrations(CustomizeDirective<GenerationContext, Pyt
             CodegenUtils.runCommand("python3 " + file, fileManifest.getBaseDir());
         }
         formatCode(fileManifest);
+        formatDocstrings(fileManifest);
         runMypy(fileManifest);
     }
 
@@ -352,6 +353,17 @@ private void formatCode(FileManifest fileManifest) {
         CodegenUtils.runCommand("python3 -m black . --exclude \"\"", fileManifest.getBaseDir());
     }
 
+    private void formatDocstrings(FileManifest fileManifest) {
+        try {
+            CodegenUtils.runCommand("python3 -m docformatter -h", fileManifest.getBaseDir());
+        } catch (CodegenException e) {
+            LOGGER.warning("Unable to find the python package docformatter. Skipping formatting.");
+            return;
+        }
+        LOGGER.info("Running docformatter on generated code");
+        CodegenUtils.runCommand("python3 -m docformatter --recursive .", fileManifest.getBaseDir());
+    }
+
     private void runMypy(FileManifest fileManifest) {
         try {
             CodegenUtils.runCommand("python3 -m mypy -h", fileManifest.getBaseDir());
diff --git a/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithypython/localservice/extensions/DafnyPythonLocalServiceStructureGenerator.java b/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithypython/localservice/extensions/DafnyPythonLocalServiceStructureGenerator.java
index b05cdd8fb..1730b1de5 100644
--- a/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithypython/localservice/extensions/DafnyPythonLocalServiceStructureGenerator.java
+++ b/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithypython/localservice/extensions/DafnyPythonLocalServiceStructureGenerator.java
@@ -374,13 +374,6 @@ protected void writeFromDict(boolean isError) {
       () -> {
         writer.writeDocs(() -> {
           writer.write("Creates a $L from a dictionary.\n", shapeName);
-          writer.write(
-            writer.formatDocs(
-              """
-              The dictionary is expected to use the modeled shape names rather \
-              than the parameter names as keys to be mostly compatible with boto3."""
-            )
-          );
         });
 
         if (shape.members().isEmpty() && !isError) {
@@ -523,13 +516,6 @@ protected void writeAsDict(boolean isError) {
             "Converts the $L to a dictionary.\n",
             symbolProvider.toSymbol(shape).getName()
           );
-          writer.write(
-            writer.formatDocs(
-              """
-              The dictionary uses the modeled shape names rather than the parameter names \
-              as keys to be mostly compatible with boto3."""
-            )
-          );
         });
 
         // If there aren't any optional members, it's best to return immediately.