2424import uk .co .real_logic .sbe .ir .Token ;
2525import org .agrona .LangUtil ;
2626import org .agrona .Verify ;
27- import org .agrona .generation .OutputManager ;
27+ import org .agrona .generation .DynamicPackageOutputManager ;
2828
2929import java .io .IOException ;
3030import java .io .Writer ;
3131import java .util .ArrayList ;
32+ import java .util .HashSet ;
3233import java .util .List ;
34+ import java .util .Set ;
3335import java .util .function .Predicate ;
3436
3537import static uk .co .real_logic .sbe .generation .Generators .toLowerFirstChar ;
@@ -49,23 +51,48 @@ public class JavaDtoGenerator implements CodeGenerator
4951 private static final String BASE_INDENT = "" ;
5052
5153 private final Ir ir ;
52- private final OutputManager outputManager ;
54+ private final DynamicPackageOutputManager outputManager ;
55+ private final boolean shouldSupportTypesPackageNames ;
56+ private final Set <String > packageNameByTypes = new HashSet <>();
5357
5458 /**
55- * Create a new C# DTO {@link CodeGenerator}.
59+ * Create a new Java DTO {@link CodeGenerator}.
5660 *
57- * @param ir for the messages and types.
58- * @param outputManager for generating the DTOs to.
61+ * @param ir for the messages and types.
62+ * @param shouldSupportTypesPackageNames generator support for types in their own package.
63+ * @param outputManager for generating the DTOs to.
5964 */
60- public JavaDtoGenerator (final Ir ir , final OutputManager outputManager )
65+ public JavaDtoGenerator (
66+ final Ir ir ,
67+ final boolean shouldSupportTypesPackageNames ,
68+ final DynamicPackageOutputManager outputManager )
6169 {
70+ this .shouldSupportTypesPackageNames = shouldSupportTypesPackageNames ;
6271 Verify .notNull (ir , "ir" );
6372 Verify .notNull (outputManager , "outputManager" );
6473
6574 this .ir = ir ;
6675 this .outputManager = outputManager ;
6776 }
6877
78+ /**
79+ * Fetch the type's explicit package - if set and should be supported.
80+ *
81+ * @param token the 0-th token of the type.
82+ * @param ir the intermediate representation.
83+ * @return the overridden package name of the type if set and supported, or {@link Ir#applicableNamespace()}.
84+ */
85+ private String fetchTypesPackageName (final Token token , final Ir ir )
86+ {
87+ if (shouldSupportTypesPackageNames && token .packageName () != null )
88+ {
89+ outputManager .setPackageName (token .packageName ());
90+ return token .packageName ();
91+ }
92+
93+ return ir .applicableNamespace ();
94+ }
95+
6996 /**
7097 * {@inheritDoc}
7198 */
@@ -110,9 +137,10 @@ public void generate() throws IOException
110137 generateDisplay (classBuilder , encoderClassName , "computeEncodedLength()" ,
111138 BASE_INDENT + INDENT );
112139
140+ final String packageName = fetchTypesPackageName (msgToken , ir );
113141 try (Writer out = outputManager .createOutput (dtoClassName ))
114142 {
115- out .append (generateDtoFileHeader (ir . applicableNamespace () ));
143+ out .append (generateDtoFileHeader (packageName ));
116144 out .append ("import org.agrona.DirectBuffer;\n " );
117145 out .append ("import org.agrona.MutableDirectBuffer;\n " );
118146 out .append ("import org.agrona.concurrent.UnsafeBuffer;\n \n " );
@@ -1582,6 +1610,19 @@ private static String formatDtoClassName(final String name)
15821610
15831611 private void generateDtosForTypes () throws IOException
15841612 {
1613+ if (shouldSupportTypesPackageNames )
1614+ {
1615+ for (final List <Token > tokens : ir .types ())
1616+ {
1617+ final String packageName = tokens .get (0 ).packageName ();
1618+
1619+ if (packageName != null )
1620+ {
1621+ packageNameByTypes .add (packageName );
1622+ }
1623+ }
1624+ }
1625+
15851626 for (final List <Token > tokens : ir .types ())
15861627 {
15871628 switch (tokens .get (0 ).signal ())
@@ -1607,10 +1648,11 @@ private void generateComposite(final List<Token> tokens) throws IOException
16071648 final String encoderClassName = encoderName (name );
16081649 final String decoderClassName = decoderName (name );
16091650
1651+ final String packageName = fetchTypesPackageName (tokens .get (0 ), ir );
16101652 try (Writer out = outputManager .createOutput (className ))
16111653 {
16121654 final List <Token > compositeTokens = tokens .subList (1 , tokens .size () - 1 );
1613- out .append (generateDtoFileHeader (ir . applicableNamespace () ));
1655+ out .append (generateDtoFileHeader (packageName ));
16141656 out .append ("import org.agrona.DirectBuffer;\n " );
16151657 out .append ("import org.agrona.MutableDirectBuffer;\n " );
16161658 out .append ("import org.agrona.concurrent.UnsafeBuffer;\n \n " );
@@ -1638,10 +1680,11 @@ private void generateChoiceSet(final List<Token> tokens) throws IOException
16381680 final String encoderClassName = encoderName (name );
16391681 final String decoderClassName = decoderName (name );
16401682
1683+ final String packageName = fetchTypesPackageName (tokens .get (0 ), ir );
16411684 try (Writer out = outputManager .createOutput (className ))
16421685 {
16431686 final List <Token > setTokens = tokens .subList (1 , tokens .size () - 1 );
1644- out .append (generateDtoFileHeader (ir . applicableNamespace () ));
1687+ out .append (generateDtoFileHeader (packageName ));
16451688 out .append (generateDocumentation (BASE_INDENT , tokens .get (0 )));
16461689
16471690 final ClassBuilder classBuilder = new ClassBuilder (className , BASE_INDENT , "public final" );
@@ -1795,13 +1838,35 @@ private void generateCompositePropertyElements(
17951838 }
17961839 }
17971840
1798- private static CharSequence generateDtoFileHeader (final String packageName )
1841+ private static StringBuilder generateImportStatements (final Set <String > packages , final String currentPackage )
1842+ {
1843+ final StringBuilder importStatements = new StringBuilder ();
1844+
1845+ for (final String candidatePackage : packages )
1846+ {
1847+ if (!candidatePackage .equals (currentPackage ))
1848+ {
1849+ importStatements .append ("import " ).append (candidatePackage ).append (".*;\n " );
1850+ }
1851+ }
1852+
1853+ if (!importStatements .isEmpty ())
1854+ {
1855+ importStatements .append ("\n \n " );
1856+ }
1857+
1858+ return importStatements ;
1859+ }
1860+
1861+ private CharSequence generateDtoFileHeader (final String packageName )
17991862 {
18001863 final StringBuilder sb = new StringBuilder ();
18011864
18021865 sb .append ("/* Generated SBE (Simple Binary Encoding) message DTO */\n " );
18031866 sb .append ("package " ).append (packageName ).append (";\n \n " );
18041867
1868+ sb .append (generateImportStatements (packageNameByTypes , packageName ));
1869+
18051870 return sb ;
18061871 }
18071872
0 commit comments