Merge pull request #27 from Medium/nick-extend

Add extends to pbnj. This will let us break cycles by letting us add fields later.

lib/Project.js

@@ -66,6 +66,9 @@ function Project(opt_basePath) {
    * @private {Array.<{proto: string, template: string, suffix: string}>}
    */
   this._compileJobs = []
+
+  /** @private {boolean} */
+  this._resolvedExtensions = false
 }
 module.exports = Project
 
@@ -201,11 +204,53 @@ Project.prototype._processProto = function (fileName) {
 }
 
 
+/**
+ * Resolve all protobuf extensions into the main messages.
+ * @private
+ */
+Project.prototype._resolveExtensions = function () {
+  if (this._resolvedExtensions) return
+
+  var protos = this.getProtos()
+  var messagesByName = {}
+  var extensionsByName = {}
+  var i, j, descriptor;
+
+  for (i = 0; i < protos.length; i++) {
+    descriptor = protos[i]
+    var messages = descriptor.getMessages()
+    for (j = 0; j < messages.length; j++) {
+      var m = messages[j]
+      var mName = descriptor.getPackage() + '.' + m.getName()
+
+      messagesByName[mName] = m
+    }
+  }
+
+  for (i = 0; i < protos.length; i++) {
+    descriptor = protos[i]
+    var extensions = descriptor.getExtends()
+    for (j = 0; j < extensions.length; j++) {
+      var e = extensions[j]
+      var eName = descriptor.getPackage() + '.' + e.getName()
+      if (messagesByName[eName]) {
+        e.mergeInto(messagesByName[eName])
+      }
+    }
+  }
+
+  this._resolvedExtensions = true
+}
+
+
+
 /**
  * Executes all the compilation jobs.
  * @return {Promise} A promise of when all compilation jobs have finished
  */
 Project.prototype.compile = function () {
+  this._resolveExtensions()
+
   soynode.setOptions({
     outputDir: this._outDir,
     uniqueDir: false,

lib/descriptors/ExtendDescriptor.js

@@ -0,0 +1,88 @@
+
+/**
+ * @fileoverview Extend class for extending a pb message.
+ */
+
+var util = require('util')
+var Descriptor = require('./Descriptor')
+var helper = require('../helper')
+
+
+/**
+ * @param {string} name The message name
+ * @constructor
+ * @extends {Descriptor}
+ */
+function ExtendDescriptor(name) {
+  Descriptor.call(this)
+
+  this._name = name
+  this._fields = {}
+  this._fieldNames = {}
+}
+util.inherits(ExtendDescriptor, Descriptor)
+module.exports = ExtendDescriptor
+
+
+/**
+ * Merge this Extend into the original message.
+ */
+ExtendDescriptor.prototype.mergeInto = function (message) {
+  for (var key in this._fields) {
+    var field = this._fields[key]
+    if (message.getFieldByTag(field.getTag())) {
+      throw new Error(
+          'duplicate tag name in extend "' + this._name + '" at field ' + field.getName())
+    }
+
+    message.addField(field)
+  }
+}
+
+
+/** @override */
+ExtendDescriptor.prototype.inspect = function (depth) {
+  return util.inspect(this.toTemplateObject(), false, null)
+}
+
+
+/** @override */
+ExtendDescriptor.prototype.toTemplateObject = function () {
+  return {
+    name: this._name,
+    fields: helper.values(this._fields, helper.toTemplateObject)
+  }
+}
+
+
+ExtendDescriptor.prototype.getName = function () {
+  return this._name
+}
+
+
+ExtendDescriptor.prototype.addField = function (field) {
+  this._fields[field.getTag()] = field
+  this._fieldNames[field.getName()] = field
+  field.setParent(this)
+  return this
+}
+
+
+ExtendDescriptor.prototype.getField = function (name) {
+  return this._fieldNames[name] || this._fieldNames[helper.toProtoCase(name)] || null
+}
+
+
+ExtendDescriptor.prototype.getFieldByTag = function (tag) {
+  return this._fields[tag] || null
+}
+
+
+ExtendDescriptor.prototype.getFields = function () {
+  return helper.values(this._fields)
+}
+
+
+ExtendDescriptor.prototype.getFieldNames = function () {
+  return Object.keys(this._fieldNames)
+}

lib/descriptors/ProtoDescriptor.js

@@ -23,6 +23,7 @@ function ProtoDescriptor(filePath) {
   this._imports = []
   this._messages = {}
   this._enums = {}
+  this._extends = {}
 
   this._rescursing = false
 }
@@ -36,6 +37,7 @@ ProtoDescriptor.prototype.inspect = function () {
     name: this.getName(),
     package: this._package,
     messages: this._messages,
+    extends: this._extends,
     enums: this._enums,
     options: this._options,
     importNames: this._importNames,
@@ -56,6 +58,7 @@ ProtoDescriptor.prototype.toTemplateObject = function () {
     package: this._package,
     options: this._options,
     messages: helper.values(this._messages, helper.toTemplateObject),
+    extends: helper.values(this._messages, helper.toTemplateObject),
     enums: helper.values(this._enums, helper.toTemplateObject),
     importNames: this._importNames,
     imports: this._imports.map(helper.toTemplateObject)
@@ -141,6 +144,23 @@ ProtoDescriptor.prototype.getMessageTypes = function () {
 }
 
 
+ProtoDescriptor.prototype.addExtend = function (extend) {
+  this._extends[extend.getName()] = extend
+  extend.setParent(this)
+  return this
+}
+
+
+ProtoDescriptor.prototype.getExtend = function (name) {
+  return this._extends[name]
+}
+
+
+ProtoDescriptor.prototype.getExtends = function () {
+  return helper.values(this._extends)
+}
+
+
 ProtoDescriptor.prototype.addEnum = function (enumeration) {
   this._enums[enumeration.getName()] = enumeration
   enumeration.setParent(this)

lib/parser.js

@@ -8,6 +8,7 @@ var Token = require('./Token')
 var ProtoDescriptor = require('./descriptors/ProtoDescriptor')
 var MessageDescriptor = require('./descriptors/MessageDescriptor')
 var EnumDescriptor = require('./descriptors/EnumDescriptor')
+var ExtendDescriptor = require('./descriptors/ExtendDescriptor')
 var FieldDescriptor = require('./descriptors/FieldDescriptor')
 var ParseError = require('./errors').ParseError
 var FieldType = require('./FieldType')
@@ -20,7 +21,6 @@ Object.keys(FieldType).forEach(function (k) {
 
 // TODO : validate enums
 // TODO : validate types
-// TODO : support for proto extensions ?
 // TODO : Implement experimental Map type through options.
 
 
@@ -206,8 +206,23 @@ module.exports = function parser(identifier, string) {
   // Parses extend: extend google.protobuf.MessageOptions {optional string type = 50001;}
   function parseExtend(parent) {
     var name = expect(Token.Type.WORD).content
+    var extend = new ExtendDescriptor(name)
+    parent.addExtend(extend)
+
     parseBlock(function (token) {
-      // TODO : parse and add to descriptor
+      switch (token.content) {
+        case 'required':
+          parseRequiredField(extend)
+          break
+        case 'optional':
+          parseOptionalField(extend)
+          break
+        case 'repeated':
+          parseRepeatedField(extend)
+          break
+        default:
+          throw new ParseError(identifier, token)
+      }
     })
   }
 

tests/extensions_test.js

@@ -0,0 +1,37 @@
+
+/**
+ * @fileoverview Tests that extensions work.
+ */
+
+var fs = require('fs')
+var parser = require('../lib/parser')
+var path = require('path')
+var Project = require('../lib/Project')
+var baseDir = __dirname
+
+exports.testExtendsParsing = function (test) {
+  var proto = parseFile('super-person.proto')
+  test.equal(1, proto.getExtends().length)
+
+  var e = proto.getExtend('Person')
+  test.ok(e)
+
+  test.equal('number', e.getField('flying_speed').getType())
+  test.done()
+}
+
+exports.testExtendsResolution = function (test) {
+  var project = new Project(baseDir)
+      .addProto('protos/person.proto')
+      .addProto('protos/super-person.proto')
+
+  project._resolveExtensions()
+  var person = project.getProtos('protos/person.proto')[0]
+  var personMsg = person.getMessage('Person')
+  test.equal('number', personMsg.getField('flying_speed').getType())
+  test.done()
+}
+
+function parseFile(file) {
+  return parser(file, fs.readFileSync(path.join(__dirname, 'protos', file), 'utf8'))
+}

tests/protos/person.proto

@@ -29,4 +29,6 @@ message Person {
     optional PhoneType type = 2 [default = HOME];
     optional string countryCode = 3 [default = "+1"];
   }
+
+  extensions 100 to 199;
 }

tests/protos/super-person.proto

@@ -0,0 +1,11 @@
+/**
+ * Extension of person
+ */
+
+import "protos/person.proto";
+
+package examples;
+
+extend Person {
+  optional int32 flying_speed = 101;
+}