From 59b2d90361db9f7549ebdfa6dfb059b330cbe14a Mon Sep 17 00:00:00 2001
From: xxyzz <gitpull@protonmail.com>
Date: Tue, 8 Dec 2020 11:22:18 +0800
Subject: [PATCH] add option to enable lemmatization

---
 __init__.py |  2 +-
 config.py   | 13 +++++++++++--
 database.py |  9 ++++++++-
 ui.py       | 11 +++++------
 4 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/__init__.py b/__init__.py
index 71a7e9c..c2b94cc 100644
--- a/__init__.py
+++ b/__init__.py
@@ -19,4 +19,4 @@ def config_widget(self):
         return ConfigWidget()
 
     def save_settings(self, config_widget):
-        pass
+        config_widget.save_settings()
diff --git a/config.py b/config.py
index 0ac2269..5ce32db 100644
--- a/config.py
+++ b/config.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
-from calibre.utils.config import JSONConfig
-from PyQt5.Qt import QWidget, QPushButton, QVBoxLayout
 import webbrowser
 
+from calibre.utils.config import JSONConfig
+from PyQt5.Qt import QPushButton, QRadioButton, QVBoxLayout, QWidget
+
 prefs = JSONConfig('plugins/worddumb')
+prefs.defaults['lemmatize'] = True
 
 
 class ConfigWidget(QWidget):
@@ -13,6 +15,10 @@ def __init__(self):
         self.vl = QVBoxLayout()
         self.setLayout(self.vl)
 
+        self.lemmatize_button = QRadioButton('Lemmatize', self)
+        self.lemmatize_button.setChecked(prefs['lemmatize'])
+        self.vl.addWidget(self.lemmatize_button)
+
         self.donate_button = QPushButton('Donate', self)
         self.donate_button.clicked.connect(self.donate)
         self.vl.addWidget(self.donate_button)
@@ -26,3 +32,6 @@ def donate(self):
 
     def github(self):
         webbrowser.open('https://github.com/xxyzz/WordDumb')
+
+    def save_settings(self):
+        prefs['lemmatize'] = self.lemmatize_button.isChecked()
diff --git a/database.py b/database.py
index 6f6b48c..a25eaa1 100644
--- a/database.py
+++ b/database.py
@@ -2,6 +2,8 @@
 import sqlite3
 from pathlib import Path
 
+from calibre_plugins.worddumb.config import prefs
+
 
 def connect_ww_database():
     ww_conn = sqlite3.connect(":memory:")
@@ -59,8 +61,13 @@ def create_lang_layer(asin, book_path):
 
 
 def match_lemma(start, word, ll_conn, ww_conn):
+    word = word.lower()
+    if prefs['lemmatize']:
+        from nltk.corpus import wordnet as wn
+        word = wn.morphy(word)
+
     for result in ww_conn.execute("SELECT * FROM words WHERE lemma = ?",
-                                  (word.lower(), )):
+                                  (word, )):
         (_, sense_id, difficulty) = result
         ll_conn.execute('''
         INSERT INTO glosses (start, difficulty, sense_id, low_confidence)
diff --git a/ui.py b/ui.py
index f8c5f96..9383a9f 100644
--- a/ui.py
+++ b/ui.py
@@ -24,12 +24,11 @@ def run(self):
 
     def install_libs(self):
         extarct_path = Path(config_dir).joinpath('plugins/worddumb')
-        if extarct_path.is_dir():
-            return
-        with ZipFile(self.plugin_path, 'r') as zf:
-            for f in zf.namelist():
-                if '.venv' in f:
-                    zf.extract(f, path=extarct_path)
+        if not extarct_path.is_dir():
+            with ZipFile(self.plugin_path, 'r') as zf:
+                for f in zf.namelist():
+                    if '.venv' in f:
+                        zf.extract(f, path=extarct_path)
         for dir in extarct_path.joinpath('.venv/lib').iterdir():
             sys.path.append(str(dir.joinpath('site-packages')))
         import nltk