com.squareup.wire.sample.CodegenSample Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of wire-codegen-sample Show documentation
There is a newer version: 3.1.0
/*
 * Copyright (C) 2016 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.wire.sample;

import com.google.common.base.Stopwatch;
import com.google.common.collect.ImmutableSet;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.TypeSpec;
import com.squareup.wire.java.JavaGenerator;
import com.squareup.wire.schema.EnumType;
import com.squareup.wire.schema.IdentifierSet;
import com.squareup.wire.schema.Location;
import com.squareup.wire.schema.MessageType;
import com.squareup.wire.schema.ProtoFile;
import com.squareup.wire.schema.Schema;
import com.squareup.wire.schema.SchemaLoader;
import com.squareup.wire.schema.Service;
import com.squareup.wire.schema.Type;
import java.io.File;
import java.io.IOException;

/**
 * Reads in .proto files, links ’em together, prunes ’em, and emits both types and services. This
 * can be used to either complement or replace {@code WireCompiler}.
 */
public final class CodegenSample {
  private final Log log;
  private final ImmutableSet sources;
  private final ImmutableSet protos;
  private final String generatedSourceDirectory;
  private final IdentifierSet identifierSet;

  public CodegenSample(Log log, ImmutableSet sources,
      ImmutableSet protos, String generatedSourceDirectory,
      IdentifierSet identifierSet) {
    this.log = log;
    this.sources = sources;
    this.protos = protos;
    this.generatedSourceDirectory = generatedSourceDirectory;
    this.identifierSet = identifierSet;
  }

  public void execute() throws IOException {
    Schema schema = loadSchema();

    if (!identifierSet.isEmpty()) {
      schema = retainRoots(schema);
    }

    JavaGenerator javaGenerator = JavaGenerator.get(schema);
    ServiceGenerator serviceGenerator = new ServiceGenerator(javaGenerator);

    for (ProtoFile protoFile : schema.protoFiles()) {
      for (Type type : protoFile.types()) {
        Stopwatch stopwatch = Stopwatch.createStarted();
        ClassName javaTypeName = (ClassName) javaGenerator.typeName(type.type());
        TypeSpec typeSpec = type instanceof MessageType
            ? javaGenerator.generateMessage((MessageType) type)
            : javaGenerator.generateEnum((EnumType) type);
        writeJavaFile(javaTypeName, typeSpec, type.location(), stopwatch);
      }

      for (Service service : protoFile.services()) {
        Stopwatch stopwatch = Stopwatch.createStarted();
        ClassName javaTypeName = (ClassName) javaGenerator.typeName(service.type());
        TypeSpec typeSpec = serviceGenerator.api(service);
        writeJavaFile(javaTypeName, typeSpec, service.location(), stopwatch);
      }
    }
  }

  private Schema loadSchema() throws IOException {
    Stopwatch stopwatch = Stopwatch.createStarted();
    SchemaLoader loader = new SchemaLoader();
    for (String source : sources) {
      loader.addSource(new File(source));
    }
    for (String proto : protos) {
      loader.addProto(proto);
    }
    Schema schema = loader.load();
    log.info("Loaded %s proto files in %s", schema.protoFiles().size(), stopwatch);
    return schema;
  }

  private Schema retainRoots(Schema schema) {
    Stopwatch stopwatch = Stopwatch.createStarted();
    int oldSize = countTypes(schema);

    Schema prunedSchema = schema.prune(identifierSet);
    int newSize = countTypes(prunedSchema);

    log.info("Pruned schema from %s types to %s types in %s", oldSize, newSize, stopwatch);

    return prunedSchema;
  }

  private void writeJavaFile(ClassName javaTypeName, TypeSpec typeSpec, Location location,
      Stopwatch stopwatch) throws IOException {
    JavaFile.Builder builder = JavaFile.builder(javaTypeName.packageName(), typeSpec)
        .addFileComment("Code generated by $L, do not edit.", CodegenSample.class.getName());
    if (location != null) {
      builder.addFileComment("\nSource file: $L", location.path());
    }
    JavaFile javaFile = builder.build();
    try {
      javaFile.writeTo(new File(generatedSourceDirectory));
    } catch (IOException e) {
      throw new IOException("Failed to write " + javaFile.packageName + "."
          + javaFile.typeSpec.name + " to " + generatedSourceDirectory, e);
    }
    log.info("Generated %s in %s", javaTypeName, stopwatch);
  }

  private int countTypes(Schema prunedSchema) {
    int result = 0;
    for (ProtoFile protoFile : prunedSchema.protoFiles()) {
      result += protoFile.types().size();
    }
    return result;
  }

  interface Log {
    void info(String format, Object... args);
  }
}