From 96f2955e94c58ec752247b2614872c1236c6655f Mon Sep 17 00:00:00 2001
From: Tobe O <thosakwe@gmail.com>
Date: Sat, 14 Jul 2018 18:06:08 -0400
Subject: [PATCH] Update custom service generator

---
 lib/src/commands/service_generators/custom.dart | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/lib/src/commands/service_generators/custom.dart b/lib/src/commands/service_generators/custom.dart
index 7d1f7f9..a1de5f3 100644
--- a/lib/src/commands/service_generators/custom.dart
+++ b/lib/src/commands/service_generators/custom.dart
@@ -12,13 +12,16 @@ class CustomServiceGenerator extends ServiceGenerator {
 
   @override
   void applyToLibrary(LibraryBuilder library, String name, String lower) {
-    var clazz = new ClassBuilder('${name}Service', asExtends: new TypeBuilder('Service'));
-    library.addMember(clazz);
+    library.body.add(new Class((clazz) {
+      clazz
+        ..name = '${name}Service'
+        ..extend = refer('Service');
+    }));
   }
 
   @override
-  ExpressionBuilder createInstance(
+  Expression createInstance(
       MethodBuilder methodBuilder, String name, String lower) {
-    return new TypeBuilder('${name}Service').newInstance([]);
+    return refer('${name}Service').newInstance([]);
   }
 }