60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
import os
|
|
import pathlib
|
|
import pickle
|
|
import sqlite3
|
|
import tempfile
|
|
|
|
from tweets import Tweet
|
|
|
|
class State:
|
|
def __init__(self, db: str) -> None:
|
|
self._file = db
|
|
con = self._connect()
|
|
cur = con.cursor()
|
|
cur.execute('''CREATE TABLE IF NOT EXISTS seen_tweets (
|
|
id INTEGER PRIMARY KEY,
|
|
tweet_id TEXT UNIQUE);''')
|
|
con.commit()
|
|
con.close()
|
|
|
|
|
|
def _connect(self):
|
|
return sqlite3.connect(self._file)
|
|
|
|
def add(self, tweet: Tweet):
|
|
con = self._connect()
|
|
cur = con.cursor()
|
|
cur.execute("INSERT INTO seen_tweets(tweet_id) VALUES (?)", (tweet.id,))
|
|
con.commit()
|
|
|
|
def has(self, tweet: Tweet):
|
|
con = self._connect()
|
|
cur = con.cursor()
|
|
rows = cur.execute("SELECT * FROM seen_tweets WHERE tweet_id=?", (tweet.id,))
|
|
return rows.fetchone() is not None
|
|
|
|
|
|
if __name__ == '__main__':
|
|
with tempfile.NamedTemporaryFile() as tmp:
|
|
state = State(tmp.name)
|
|
|
|
print('Inserting t1, t2')
|
|
t1 = Tweet(id="test1", author="test1", content="test1")
|
|
t2 = Tweet(id="test2", author="test2", content="test2")
|
|
t3 = Tweet(id="test3", author="test3", content="test3")
|
|
state.add(t1)
|
|
state.add(t2)
|
|
|
|
print('Checking if they exist in db')
|
|
print('t1:', state.has(t1))
|
|
assert state.has(t1)
|
|
print('t2:', state.has(t2))
|
|
assert state.has(t2)
|
|
print('t3:', state.has(t3))
|
|
assert not state.has(t3)
|
|
|
|
|
|
|
|
|
|
|