Allow services to be generated in their own files

1 files changed, 74 insertions, 41 deletions

diff --git a/nimpb/compiler/generator.nim b/nimpb/compiler/generator.nim
index fde4da0..8e144c0 100644
--- a/nimpb/compiler/generator.nim
+++ b/nimpb/compiler/generator.nim
@@ -61,7 +61,22 @@ type
         Proto3
 
     ServiceGenerator* = ref object of RootObj
+        ## If set, the service will be generated in it's own file. The filename
+        ## where the service will be generated will be <basename>_<suffix>.nim,
+        ## where basename is the name of the proto file without the .proto
+        ## suffix.
+        fileSuffix*: string
+
+        ## This will be set to the basename (without extension) of the file
+        ## where the serialization code will be generated. It will be set before
+        ## any callbacks are called. You can use this for importing if you
+        ## generate a new file for the service by setting fileSuffix.
+        fileName*: string
+
+        ## This will be called once per proto file.
         genImports*: proc (): string
+
+        ## This will be called once per service definition.
         genService*: proc (service: Service): string
 
     Service* = ref object
@@ -978,17 +993,37 @@ iterator genProcs(msg: Message): string =
         yield indent(&"result = read{msg.names}(ss)", 4)
         yield ""
 
+proc hasGenImports(serviceGenerator: ServiceGenerator): bool =
+    serviceGenerator != nil and serviceGenerator.genImports != nil
+
+proc hasGenService(serviceGenerator: ServiceGenerator): bool =
+    serviceGenerator != nil and serviceGenerator.genService != nil
+
+proc ownFile(serviceGenerator: ServiceGenerator): bool =
+    serviceGenerator != nil and serviceGenerator.fileSuffix != nil
+
 proc processFile(fdesc: google_protobuf_FileDescriptorProto,
                  otherFiles: TableRef[string, ProtoFile],
-                 serviceGenerator: ServiceGenerator): ProcessedFile =
+                 serviceGenerator: ServiceGenerator): seq[ProcessedFile] =
+    result = @[]
+
     var (dir, name, _) = splitFile(fdesc.name)
     var pbfilename = (dir / name) & "_pb.nim"
 
     log(&"processing {fdesc.name}: {pbfilename}")
 
-    new(result)
-    result.name = pbfilename
-    result.data = ""
+    var pbFile = ProcessedFile(name: pbfilename, data: "")
+    var serviceFile = pbFile
+
+    add(result, pbFile)
+
+    if serviceGenerator != nil:
+        serviceGenerator.fileName = name & "_pb"
+
+    if ownFile(serviceGenerator):
+        var serviceFilename = (dir / name) & "_" & serviceGenerator.fileSuffix & ".nim"
+        serviceFile = ProcessedFile(name: serviceFilename, data: "")
+        add(result, serviceFile)
 
     let parsed = parseFile(fdesc.name, fdesc)
 
@@ -1002,22 +1037,25 @@ proc processFile(fdesc: google_protobuf_FileDescriptorProto,
         if tmp:
             hasMaps = true
 
-    addLine(result.data, "# Generated by protoc_gen_nim. Do not edit!")
-    addLine(result.data, "")
-    addLine(result.data, "import base64")
-    addLine(result.data, "import intsets")
-    addLine(result.data, "import json")
+    addLine(pbFile.data, "# Generated by protoc_gen_nim. Do not edit!")
+    addLine(pbFile.data, "")
+    addLine(pbFile.data, "import base64")
+    addLine(pbFile.data, "import intsets")
+    addLine(pbFile.data, "import json")
     if hasMaps:
-        addLine(result.data, "import tables")
-        addLine(result.data, "export tables")
-    addLine(result.data, "")
-    addLine(result.data, "import nimpb/nimpb")
-    addLine(result.data, "import nimpb/json as nimpb_json")
-    addLine(result.data, "")
+        addLine(pbFile.data, "import tables")
+        addLine(pbFile.data, "export tables")
+    addLine(pbFile.data, "")
+    addLine(pbFile.data, "import nimpb/nimpb")
+    addLine(pbFile.data, "import nimpb/json as nimpb_json")
+    addLine(pbFile.data, "")
 
-    if serviceGenerator != nil:
-        if serviceGenerator.genImports != nil:
-            add(result.data, serviceGenerator.genImports())
+    if hasGenImports(serviceGenerator):
+        add(serviceFile.data, serviceGenerator.genImports())
+
+    # if serviceGenerator != nil:
+    #     if serviceGenerator.genImports != nil:
+    #         add(pbFile.data, serviceGenerator.genImports())
 
     for dep in fdesc.dependency:
         var (dir, depname, _) = splitFile(dep)
@@ -1026,37 +1064,35 @@ proc processFile(fdesc: google_protobuf_FileDescriptorProto,
             dir = "nimpb/wkt"
 
         var deppbname = (dir / depname) & "_pb"
-        addLine(result.data, &"import {deppbname}")
+        addLine(pbFile.data, &"import {deppbname}")
 
     if hasDependency(fdesc):
-        addLine(result.data, "")
+        addLine(pbFile.data, "")
 
-    addLine(result.data, "type")
+    addLine(pbFile.data, "type")
 
     for e in parsed.enums:
-        for line in genType(e): addLine(result.data, indent(line, 4))
+        for line in genType(e): addLine(pbFile.data, indent(line, 4))
 
     for message in parsed.messages:
-        for line in genType(message): addLine(result.data, indent(line, 4))
+        for line in genType(message): addLine(pbFile.data, indent(line, 4))
 
-    addLine(result.data, "")
+    addLine(pbFile.data, "")
 
     for message in sortDependencies(parsed.messages):
         for line in genMessageProcForwards(message):
-            addLine(result.data, line)
-        addLine(result.data, "")
+            addLine(pbFile.data, line)
+        addLine(pbFile.data, "")
 
     for message in sortDependencies(parsed.messages):
         for line in genProcs(message):
-            addLine(result.data, line)
-        addLine(result.data, "")
+            addLine(pbFile.data, line)
+        addLine(pbFile.data, "")
 
-    if serviceGenerator != nil:
-        if serviceGenerator.genService != nil:
-            for serviceDesc in fdesc.service:
-                let service = newService(serviceDesc, parsed)
-                addLine(result.data, "")
-                add(result.data, serviceGenerator.genService(service))
+    if hasGenService(serviceGenerator):
+        for serviceDesc in fdesc.service:
+            let service = newService(serviceDesc, parsed)
+            add(serviceFile.data, serviceGenerator.genService(service))
 
 proc processFileDescriptorSet*(filename: string,
                                outdir: string,
@@ -1079,10 +1115,7 @@ proc processFileDescriptorSet*(filename: string,
 
     for file in fileSet.file:
         if (file.name in protos) or ((basePath / file.name) in protos):
-            let processedFile = processFile(file, otherFiles, serviceGenerator)
-
-            let fullPath = outdir / processedFile.name
-
-            createDir(parentDir(fullPath))
-
-            writeFile(fullPath, processedFile.data)
+            for processedFile in processFile(file, otherFiles, serviceGenerator):
+                let fullPath = outdir / processedFile.name
+                createDir(parentDir(fullPath))
+                writeFile(fullPath, processedFile.data)