From bc58dae3a63a38a84f1eeec87122e72b5c2a2185 Mon Sep 17 00:00:00 2001
From: Kevin Moore <kevmoo@google.com>
Date: Thu, 9 Aug 2018 13:20:10 -0700
Subject: [PATCH] fix(LibraryReader) `classes` should return all classes,
 including parts

Also prepare to release 0.9.0+1
---
 source_gen/CHANGELOG.md                     |  5 +++++
 source_gen/lib/src/library.dart             |  2 +-
 source_gen/pubspec.yaml                     |  2 +-
 source_gen/test/library/find_type_test.dart | 20 ++++++++++++++++++--
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/source_gen/CHANGELOG.md b/source_gen/CHANGELOG.md
index c423f04b..ea2fd651 100644
--- a/source_gen/CHANGELOG.md
+++ b/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
diff --git a/source_gen/lib/src/library.dart b/source_gen/lib/src/library.dart
index 8104b7db..238f5b8a 100644
--- a/source_gen/lib/src/library.dart
+++ b/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) {
diff --git a/source_gen/pubspec.yaml b/source_gen/pubspec.yaml
index e39486ee..e79b9c17 100644
--- a/source_gen/pubspec.yaml
+++ b/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
diff --git a/source_gen/test/library/find_type_test.dart b/source_gen/test/library/find_type_test.dart
index cd7812e5..7ffbf8d4 100644
--- a/source_gen/test/library/find_type_test.dart
+++ b/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);
   });