fix(LibraryReader) classes should return all classes, including parts

source_gen/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.9.0+1
+
+* Fix `LibraryReader.classElements` to return classes from parts, if they exist,
+  as well as from the defining library file.
+
 ## 0.9.0
 
 * Introduce `SharedPartBuilder` for creating part files that can be merged

source_gen/lib/src/library.dart

@@ -166,7 +166,7 @@ class LibraryReader {
 
   /// All of the `class` elements in this library.
   Iterable<ClassElement> get classElements =>
-      element.definingCompilationUnit.types;
+      element.units.expand((cu) => cu.types);
 
   static Iterable<Element> _getElements(CompilationUnitMember member) {
     if (member is TopLevelVariableDeclaration) {

source_gen/pubspec.yaml

@@ -1,5 +1,5 @@
 name: source_gen
-version: 0.9.0
+version: 0.9.0+1
 author: Dart Team <misc@dartlang.org>
 description: Automated source code generation for Dart.
 homepage: https://github.com/dart-lang/source_gen

source_gen/test/library/find_type_test.dart

@@ -14,23 +14,39 @@ final _source = r'''
   export 'package:source_gen/source_gen.dart' show Generator;
   import 'dart:async' show Stream;
 
+  part 'part.dart';
+
   class Example {}
 ''';
 
+final _partSource = r'''
+part of 'source.dart';
+
+class PartClass {}
+''';
+
 void main() {
   LibraryReader library;
 
   setUpAll(() async {
-    library = await resolveSource(
-      _source,
+    library = await resolveSources(
+      {'a|source.dart': _source, 'a|part.dart': _partSource},
       (r) async => new LibraryReader(await r.findLibraryByName('test_lib')),
     );
   });
 
+  test('class count', () {
+    expect(library.classElements, hasLength(2));
+  });
+
   test('should return a type not exported', () {
     expect(library.findType('Example'), _isClassElement);
   });
 
+  test('should return a type from a part', () {
+    expect(library.findType('PartClass'), _isClassElement);
+  });
+
   test('should return a type exported from dart:', () {
     expect(library.findType('LinkedHashMap'), _isClassElement);
   });