diff --git a/lib/xdrgen/compilation.rb b/lib/xdrgen/compilation.rb index 8b13cdf4f..24d260793 100644 --- a/lib/xdrgen/compilation.rb +++ b/lib/xdrgen/compilation.rb @@ -2,11 +2,12 @@ module Xdrgen class Compilation extend Memoist - def initialize(source_paths, output_dir:".", language: :ruby, namespace: nil, options: {}) + def initialize(source_paths, output_dir:".", language: :ruby, generator: nil, namespace: nil, options: {}) @source_paths = source_paths @output_dir = output_dir @namespace = namespace @language = language + @generator = generator @options = options end @@ -22,8 +23,8 @@ def initialize(source_paths, output_dir:".", language: :ruby, namespace: nil, op def compile output = Output.new(@source_paths, @output_dir) - - generator = Generators.for_language(@language).new(ast, output, @namespace, @options) + generator_class = @generator || Generators.for_language(@language) + generator = generator_class.new(ast, output, @namespace, @options) generator.generate ensure output.close diff --git a/spec/lib/xdrgen/custom_generator_spec.rb b/spec/lib/xdrgen/custom_generator_spec.rb new file mode 100644 index 000000000..611242408 --- /dev/null +++ b/spec/lib/xdrgen/custom_generator_spec.rb @@ -0,0 +1,24 @@ +require 'spec_helper' + +describe Xdrgen::Generators do + it "can accept a custom generator to generate" do + expect(TestGenerator).to receive(:new).with( + an_instance_of(Xdrgen::AST::Top), + an_instance_of(Xdrgen::Output), + "namespace", + { option: true }, + ).and_call_original + expect_any_instance_of(TestGenerator).to receive(:generate).with(no_args) + Xdrgen::Compilation.new( + generator_fixture_paths, + output_dir: "output_dir/", + generator: TestGenerator, + namespace: "namespace", + options: { option: true } + ).compile + end +end + +class TestGenerator < Xdrgen::Generators::Base + def generate; end +end