@typespec/protobuf 0.49.0-dev.2 → 0.49.0-dev.4

package/src/write.ts

@@ -5,6 +5,7 @@ import {
   matchType,
   ProtoDeclaration,
   ProtoEnumDeclaration,
+  ProtoEnumVariantDeclaration,
   ProtoFieldDeclaration,
   ProtoFile,
   ProtoMessageDeclaration,
@@ -17,21 +18,19 @@ import {
 
 // This module defines how to emit the text representation of a ProtoFile AST.
 
-/**
- * Header for the top of all emitted proto files.
- *
- * We only support Protobuf 3 syntax.
- */
-export const PROTO_HEADER = `/* Generated by Microsoft TypeSpec */
-
-syntax = "proto3";
-`;
+const PROTO_MAX_ONE_LINE_DOC_LENGTH = 80;
 
 /**
  * Write the given `file` to a string.
  */
 export function writeProtoFile(file: ProtoFile): string {
-  let result = PROTO_HEADER;
+  let result = "// Generated by Microsoft TypeSpec\n";
+
+  let docComment = collect(writeBlockDocumentationComment(file)).join("\n");
+  if (docComment.length > 0) docComment = "\n" + docComment;
+  result += docComment;
+
+  result += '\nsyntax = "proto3";\n';
 
   if (file.package) result += `\npackage ${file.package};\n`;
 
@@ -51,7 +50,7 @@ export function writeProtoFile(file: ProtoFile): string {
   if (opts.length > 0) result += "\n";
 
   for (const decl of file.declarations) {
-    result += "\n" + collect(writeDeclaration(decl)).join("\n") + "\n";
+    result += "\n" + collect(writeDeclaration(decl, 0)).join("\n") + "\n";
   }
 
   return result;
@@ -60,25 +59,28 @@ export function writeProtoFile(file: ProtoFile): string {
 /**
  * Write the given `decl` to a line iterable.
  */
-function* writeDeclaration(decl: ProtoDeclaration): Iterable<string> {
+function* writeDeclaration(decl: ProtoDeclaration, indentLevel: number): Iterable<string> {
   switch (decl.kind) {
     case "message":
-      yield* writeMessage(decl);
+      yield* writeMessage(decl, indentLevel);
       return;
     case "service":
-      yield* writeService(decl);
+      yield* writeService(decl, indentLevel);
       return;
     case "field":
-      yield writeField(decl);
+      yield* writeField(decl, indentLevel);
       return;
     case "oneof":
-      yield* writeOneOf(decl);
+      yield* writeOneOf(decl, indentLevel);
       return;
     case "enum":
-      yield* writeEnum(decl);
+      yield* writeEnum(decl, indentLevel);
+      return;
+    case "variant":
+      yield* writeVariant(decl, indentLevel);
       return;
     case "method":
-      yield writeMethod(decl);
+      yield* writeMethod(decl);
       return;
     /* c8 ignore next 5 */
     default:
@@ -90,14 +92,16 @@ function* writeDeclaration(decl: ProtoDeclaration): Iterable<string> {
 /**
  * Write the given message `decl` to a line iterable.
  */
-function* writeMessage(decl: ProtoMessageDeclaration): Iterable<string> {
+function* writeMessage(decl: ProtoMessageDeclaration, indentLevel: number): Iterable<string> {
+  yield* writeBlockDocumentationComment(decl);
+
   const head = `message ${decl.name} {`;
   const tail = "}";
 
   if (decl.declarations.length > 0 || decl.reservations?.length) {
     yield head;
     yield* indent(writeReservations(decl));
-    yield* indent(flatMap(decl.declarations, writeDeclaration));
+    yield* indent(flatMap(decl.declarations, (decl) => writeDeclaration(decl, indentLevel + 1)));
     yield tail;
   } else yield head + tail;
 }
@@ -119,49 +123,65 @@ function* writeReservations(decl: ProtoMessageDeclaration): Iterable<string> {
   }
 }
 
-function* writeService(decl: ProtoServiceDeclaration): Iterable<string> {
+function* writeService(decl: ProtoServiceDeclaration, indentLevel: number): Iterable<string> {
+  yield* writeBlockDocumentationComment(decl);
+
   const head = `service ${decl.name} {`;
   const tail = "}";
 
   if (decl.operations.length > 0) {
     yield head;
-    yield* indent(flatMap(decl.operations, writeDeclaration));
+    yield* indent(flatMap(decl.operations, (decl) => writeDeclaration(decl, indentLevel + 1)));
     yield tail;
   } else yield head + tail;
 }
 
-function writeMethod(decl: ProtoMethodDeclaration): string {
+function* writeMethod(decl: ProtoMethodDeclaration): Iterable<string> {
+  yield* writeBlockDocumentationComment(decl);
+
   const [inStream, outStream] = [
     decl.stream & StreamingMode.In,
     decl.stream & StreamingMode.Out,
   ].map((v) => (v ? "stream " : ""));
 
-  return `rpc ${decl.name}(${inStream}${writeType(decl.input)}) returns (${outStream}${writeType(
+  yield `rpc ${decl.name}(${inStream}${writeType(decl.input)}) returns (${outStream}${writeType(
     decl.returns
   )});`;
 }
 
-function* writeOneOf(decl: ProtoOneOfDeclaration): Iterable<string> {
+function* writeOneOf(decl: ProtoOneOfDeclaration, indentLevel: number): Iterable<string> {
+  yield* writeBlockDocumentationComment(decl);
+
   // OneOf declarations must have at least one element, so no need to check for declarations
   yield `oneof ${decl.name} {`;
-  yield* indent(flatMap(decl.declarations, writeDeclaration));
+  yield* indent(flatMap(decl.declarations, (decl) => writeDeclaration(decl, indentLevel + 1)));
   yield "}";
 }
 
-function* writeEnum(decl: ProtoEnumDeclaration): Iterable<string> {
+function* writeEnum(decl: ProtoEnumDeclaration, indentLevel: number): Iterable<string> {
+  yield* writeBlockDocumentationComment(decl);
+
   yield `enum ${decl.name} {`;
   if (decl.allowAlias) {
     yield "  option allow_alias = true;";
 
     if (decl.variants.length > 0) yield "";
   }
-  yield* indent(flatMap(decl.variants, ([name, idx]) => `${name} = ${idx};`));
+  yield* indent(flatMap(decl.variants, (decl) => writeDeclaration(decl, indentLevel + 1)));
   yield "}";
 }
 
-function writeField(decl: ProtoFieldDeclaration): string {
+function writeVariant(decl: ProtoEnumVariantDeclaration, indentLevel: number): Iterable<string> {
+  const output = `${decl.name} = ${decl.value};`;
+
+  return writeDocumentationCommentFlexible(decl, output, indentLevel);
+}
+
+function writeField(decl: ProtoFieldDeclaration, indentLevel: number): Iterable<string> {
   const prefix = decl.repeated ? "repeated " : "";
-  return prefix + `${writeType(decl.type)} ${decl.name} = ${decl.index};`;
+  const output = prefix + `${writeType(decl.type)} ${decl.name} = ${decl.index};`;
+
+  return writeDocumentationCommentFlexible(decl, output, indentLevel);
 }
 
 function writeType(type: ProtoType): string {
@@ -189,6 +209,37 @@ function* indent(it: Iterable<string>, depth: number = 2): Iterable<string> {
   }
 }
 
+/**
+ * Writes a block comment from the given declaration.
+ */
+function* writeBlockDocumentationComment(decl: ProtoDeclaration | ProtoFile): Iterable<string> {
+  yield* decl.doc
+    ?.trim()
+    .split("\n")
+    .map((line) => `// ${line}`) ?? [];
+}
+
+/**
+ * Writes a block comment or inline postfix comment depending on the length of the content.
+ */
+function* writeDocumentationCommentFlexible(
+  decl: ProtoDeclaration,
+  output: string,
+  indentLevel: number
+): Iterable<string> {
+  const docComment = decl.doc?.trim();
+  const docCommentIsOneLine = docComment && !docComment?.includes("\n");
+
+  const fullLength = indentLevel * 2 + output.length + 1 + (docComment?.length ?? 0);
+
+  if (docCommentIsOneLine && fullLength <= PROTO_MAX_ONE_LINE_DOC_LENGTH) {
+    yield output + " // " + decl.doc;
+  } else {
+    yield* writeBlockDocumentationComment(decl);
+    yield output;
+  }
+}
+
 /**
  * A version of flatMap that works with generic iterables.
  *