support for godot and lightning (#70)

.gitattributes

@@ -0,0 +1 @@
+*.db filter=lfs diff=lfs merge=lfs -text

buster/data/{document_embeddings_huggingface.tar.gz → documents.db}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:19453cb9ec85e644306af7dcc6fcad79cbb842d15a2087a66ddf48b5cbd9fbc9
-size 46918939
+oid sha256:b86c2b4f5a2ec410c2b9132ed62213528ba10c0dc260162f689e30ba677815f1
+size 244338688

buster/docparser.py

@@ -36,6 +36,16 @@ supported_docs = {
         "filename": "documents_huggingface.csv",
         "parser": HuggingfaceParser,
     },
+    "lightning": {
+        "base_url": "https://pytorch-lightning.readthedocs.io/en/stable/",
+        "filename": "documents_lightning.csv",
+        "parser": SphinxParser,
+    },
+    "godot": {
+        "base_url": "https://docs.godotengine.org/en/stable/",
+        "filename": "documents_godot.csv",
+        "parser": SphinxParser,
+    },
 }
 
 

buster/documents/sqlite/documents.py

@@ -1,7 +1,5 @@
 import itertools
 import sqlite3
-import warnings
-import zlib
 from pathlib import Path
 from typing import Iterable, NamedTuple
 
@@ -41,7 +39,7 @@ class DocumentsDB(DocumentsManager):
     def __init__(self, db_path: sqlite3.Connection | str):
         if isinstance(db_path, (str, Path)):
             self.db_path = db_path
-            self.conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
+            self.conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES, check_same_thread=False)
         else:
             self.db_path = None
             self.conn = db_path
@@ -143,6 +141,7 @@ class DocumentsDB(DocumentsManager):
 
         sid, vid = self.add_parse(source, (section for section, _ in sections))
         self.add_chunking(sid, vid, size, (chunks for _, chunks in sections))
+        self.conn.commit()
 
     def get_documents(self, source: str) -> pd.DataFrame:
         """Get all current documents from a given source."""

buster/notebooks/db_to_csv.ipynb

@@ -10,32 +10,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "import os\n",
     "\n",
-    "from buster.docparser import read_documents, write_documents"
+    "from buster.documents import get_documents_manager_from_extension"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Path to the database\n",
     "db_path = \"documents.db\"\n",
+    "db = get_documents_manager_from_extension(db_path)(db_path)\n",
     "\n",
     "# Source to extract\n",
-    "target = \"pytorch\"\n",
-    "df = read_documents(db_path, target)\n",
+    "target = \"lightning\"\n",
+    "df = db.get_documents(target)\n",
     "\n",
     "# If you want to save it as tar.gz\n",
     "filepath = os.path.join('buster/data/', f'document_embeddings_{target}.tar.gz')\n",
-    "write_documents(filepath, target, df)"
+    "db_df = get_documents_manager_from_extension(filepath)(filepath)\n",
+    "db_df.add(target, df)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -45,7 +54,15 @@
    "name": "python3"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
    "version": "3.10.9"
   },
   "orig_nbformat": 4,