From 034d498b319d37b0639203fa4fbb304715b3ae03 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 22 Jul 2019 17:12:54 -0700
Subject: [PATCH] Support Database(memory=True) for in-memory databases

---
 docs/python-api.rst  | 10 +++++++++-
 sqlite_utils/db.py   |  9 +++++++--
 tests/conftest.py    |  5 ++---
 tests/test_create.py |  7 +++++++
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index 9444df9..378b5d9 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -27,7 +27,7 @@ If you want to create an in-memory database, you can do so like this:
 
 .. code-block:: python
 
-    db = Database(sqlite3.connect(":memory:"))
+    db = Database(sqlite3.connect(memory=True))
 
 Tables are accessed using the indexing operator, like so:
 
@@ -37,6 +37,14 @@ Tables are accessed using the indexing operator, like so:
 
 If the table does not yet exist, it will be created the first time you attempt to insert or upsert data into it.
 
+You can also access tables using the ``.table()`` method like so:
+
+.. code-block:: python
+
+    table = db.table("my_table")
+
+Using this factory function allows you to set :ref:`python_api_table_configuration`.
+
 Listing tables
 ==============
 
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index f4829ee..9f21a3b 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -85,8 +85,13 @@ class NotFoundError(Exception):
 
 
 class Database:
-    def __init__(self, filename_or_conn):
-        if isinstance(filename_or_conn, str):
+    def __init__(self, filename_or_conn=None, memory=False):
+        assert (filename_or_conn is not None and not memory) or (
+            filename_or_conn is None and memory
+        ), "Either specify a filename_or_conn or pass memory=True"
+        if memory:
+            self.conn = sqlite3.connect(":memory:")
+        elif isinstance(filename_or_conn, str):
             self.conn = sqlite3.connect(filename_or_conn)
         elif isinstance(filename_or_conn, pathlib.Path):
             self.conn = sqlite3.connect(str(filename_or_conn))
diff --git a/tests/conftest.py b/tests/conftest.py
index c7b4077..1e1fc5e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,16 +1,15 @@
 from sqlite_utils import Database
-from sqlite_utils.utils import sqlite3
 import pytest
 
 
 @pytest.fixture
 def fresh_db():
-    return Database(sqlite3.connect(":memory:"))
+    return Database(memory=True)
 
 
 @pytest.fixture
 def existing_db():
-    database = Database(sqlite3.connect(":memory:"))
+    database = Database(memory=True)
     database.conn.executescript(
         """
         CREATE TABLE foo (text TEXT);
diff --git a/tests/test_create.py b/tests/test_create.py
index 9360395..4634191 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -727,3 +727,10 @@ def test_create_table_numpy(fresh_db):
             "np.uint8": 8,
         }
     ] == list(fresh_db["types"].rows)
+
+
+def test_cannot_provide_both_filename_and_memory():
+    with pytest.raises(
+        AssertionError, match="Either specify a filename_or_conn or pass memory=True"
+    ):
+        db = Database("/tmp/foo.db", memory=True)