Refactor

src/ai.ts

@@ -8,6 +8,7 @@ import IModule from './module';
 import MessageLike from './message-like';
 import { FriendDoc } from './friend';
 import { User } from './misskey/user';
+import getCollection from './utils/get-collection';
 const ReconnectingWebSocket = require('../node_modules/reconnecting-websocket/dist/reconnecting-websocket-cjs.js');
 
 /**
@@ -55,19 +56,13 @@ export default class 藍 {
 
 	private init = () => {
 		//#region Init DB
-		this.contexts = this.db.getCollection('contexts');
-		if (this.contexts === null) {
-			this.contexts = this.db.addCollection('contexts', {
-				indices: ['key']
-			});
-		}
+		this.contexts = getCollection(this.db, 'contexts', {
+			indices: ['key']
+		});
 
-		this.friends = this.db.getCollection('friends');
-		if (this.friends === null) {
-			this.friends = this.db.addCollection('friends', {
-				indices: ['userId']
-			});
-		}
+		this.friends = getCollection(this.db, 'friends', {
+			indices: ['userId']
+		});
 		//#endregion
 
 		// Install modules

src/modules/guessing-game/index.ts

@@ -3,6 +3,7 @@ import 藍 from '../../ai';
 import IModule from '../../module';
 import MessageLike from '../../message-like';
 import serifs from '../../serifs';
+import getCollection from '../../utils/get-collection';
 
 export default class GuessingGameModule implements IModule {
 	public name = 'guessingGame';
@@ -20,12 +21,9 @@ export default class GuessingGameModule implements IModule {
 		this.ai = ai;
 
 		//#region Init DB
-		this.guesses = this.ai.db.getCollection('guessingGame');
-		if (this.guesses === null) {
-			this.guesses = this.ai.db.addCollection('guessingGame', {
-				indices: ['userId']
-			});
-		}
+		this.guesses = getCollection(this.ai.db, 'guessingGame', {
+			indices: ['userId']
+		});
 		//#endregion
 	}
 

src/modules/keyword/index.ts

@@ -3,6 +3,7 @@ import 藍 from '../../ai';
 import IModule from '../../module';
 import config from '../../config';
 import serifs from '../../serifs';
+import getCollection from '../../utils/get-collection';
 const MeCab = require('mecab-async');
 
 function kanaToHira(str: string) {
@@ -12,8 +13,6 @@ function kanaToHira(str: string) {
 	});
 }
 
-const db = '_keyword_learnedKeywords';
-
 export default class KeywordModule implements IModule {
 	public name = 'keyword';
 
@@ -28,12 +27,9 @@ export default class KeywordModule implements IModule {
 		this.ai = ai;
 
 		//#region Init DB
-		this.learnedKeywords = this.ai.db.getCollection(db);
-		if (this.learnedKeywords === null) {
-			this.learnedKeywords = this.ai.db.addCollection(db, {
-				indices: ['keyword']
-			});
-		}
+		this.learnedKeywords = getCollection(this.ai.db, '_keyword_learnedKeywords', {
+			indices: ['userId']
+		});
 		//#endregion
 
 		this.tokenizer = new MeCab();

src/utils/get-collection.ts

@@ -0,0 +1,13 @@
+import * as loki from 'lokijs';
+
+export default function(db: loki, name: string, opts?: any): loki.Collection {
+	let collection;
+
+	collection = db.getCollection(name);
+
+	if (collection === null) {
+		collection = db.addCollection(name, opts);
+	}
+
+	return collection;
+}