Allow single compilation unit to generate multiple classes

2 files changed, 6 insertions, 4 deletions

diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
index 2a842c9..6e9945e 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaFrontend.java
@@ -14,7 +14,7 @@ public class JavaFrontend implements Frontend {
     @Override
     public List<ClassNode> readSourceDirectory(Path directoryPath) throws IOException {
         JavaReader reader = new JavaReader();
-        return ExtraLists.map(
+        return ExtraLists.flatMap(
             findJavaFiles(directoryPath),
             javaFile -> reader.readClassFromFile(directoryPath, javaFile));
     }
diff --git a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
index 4f07149..446bc31 100644
--- a/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
+++ b/src/main/java/org/zwobble/couscous/frontends/java/JavaReader.java
@@ -8,25 +8,27 @@ import org.zwobble.couscous.ast.TypeName;
 
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.List;
 import java.util.Optional;
 
+import static java.util.Arrays.asList;
 import static org.zwobble.couscous.frontends.java.JavaTypes.typeOf;
 
 public class JavaReader {
     private final JavaParser parser = new JavaParser();
     
-    public ClassNode readClassFromFile(Path root, Path sourcePath) throws IOException {
+    public List<ClassNode> readClassFromFile(Path root, Path sourcePath) throws IOException {
         CompilationUnit ast = parser.parseCompilationUnit(root, sourcePath);
         return readCompilationUnit(ast);
     }
     
-    private static ClassNode readCompilationUnit(CompilationUnit ast) {
+    private static List<ClassNode> readCompilationUnit(CompilationUnit ast) {
         String name = generateClassName(ast);
         ClassNodeBuilder classBuilder = new ClassNodeBuilder(name);
         TypeDeclaration type = (TypeDeclaration)ast.types().get(0);
         readFields(type, classBuilder);
         readMethods(type, classBuilder);
-        return classBuilder.build();
+        return asList(classBuilder.build());
     }
     
     private static void readFields(TypeDeclaration type, ClassNodeBuilder classBuilder) {