simplify _session, execute

src/cs50/_session.py

@@ -6,18 +6,12 @@
 
 class Session:
     def __init__(self, url, **engine_kwargs):
-        self._url = url
-        if _is_sqlite_url(self._url):
-            _assert_sqlite_file_exists(self._url)
+        if _is_sqlite_url(url):
+            _assert_sqlite_file_exists(url)
 
-        self._engine = _create_engine(self._url, **engine_kwargs)
-        self._is_postgres = self._engine.url.get_backend_name() in {"postgres", "postgresql"}
-        _setup_on_connect(self._engine)
-        self._session = _create_scoped_session(self._engine)
-
-
-    def is_postgres(self):
-        return self._is_postgres
+        engine = _create_engine(url, **engine_kwargs)
+        _setup_on_connect(engine)
+        self._session = _create_scoped_session(engine)
 
 
     def execute(self, statement):

src/cs50/_statement.py

@@ -24,20 +24,20 @@ def __init__(self, dialect, sql, *args, **kwargs):
     def _parse(self):
         formatted_statements = sqlparse.format(self._sql, strip_comments=True).strip()
         parsed_statements = sqlparse.parse(formatted_statements)
-        num_of_statements = len(parsed_statements)
-        if num_of_statements == 0:
+        statement_count = len(parsed_statements)
+        if statement_count == 0:
             raise RuntimeError("missing statement")
-        elif num_of_statements > 1:
+        elif statement_count > 1:
             raise RuntimeError("too many statements at once")
 
         return parsed_statements[0]
 
 
     def _parse_command(self):
         for token in self._statement:
-            if token.ttype in [sqlparse.tokens.Keyword, sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML]:
+            if _is_command_token(token):
                 token_value = token.value.upper()
-                if token_value in ["BEGIN", "DELETE", "INSERT", "SELECT", "START", "UPDATE"]:
+                if token_value in {"BEGIN", "DELETE", "INSERT", "SELECT", "START", "UPDATE"}:
                     command = token_value
                     break
         else:
@@ -49,11 +49,11 @@ def _parse_command(self):
     def _bind_params(self):
         tokens = self._tokenize()
         paramstyle, placeholders = self._parse_placeholders(tokens)
-        if paramstyle in [Paramstyle.FORMAT, Paramstyle.QMARK]:
+        if paramstyle in {Paramstyle.FORMAT, Paramstyle.QMARK}:
             tokens = self._bind_format_or_qmark(placeholders, tokens)
         elif paramstyle == Paramstyle.NUMERIC:
             tokens = self._bind_numeric(placeholders, tokens)
-        if paramstyle in [Paramstyle.NAMED, Paramstyle.PYFORMAT]:
+        if paramstyle in {Paramstyle.NAMED, Paramstyle.PYFORMAT}:
             tokens = self._bind_named_or_pyformat(placeholders, tokens)
 
         tokens = _escape_verbatim_colons(tokens)
@@ -154,10 +154,10 @@ def _escape(self, value):
                 sqlalchemy.types.Boolean().literal_processor(self._dialect)(value))
 
         if isinstance(value, bytes):
-            if self._dialect.name in ["mysql", "sqlite"]:
+            if self._dialect.name in {"mysql", "sqlite"}:
                 # https://dev.mysql.com/doc/refman/8.0/en/hexadecimal-literals.html
                 return sqlparse.sql.Token(sqlparse.tokens.Other, f"x'{value.hex()}'")
-            if self._dialect.name in ["postgres", "postgresql"]:
+            if self._dialect.name in {"postgres", "postgresql"}:
                 # https://dba.stackexchange.com/a/203359
                 return sqlparse.sql.Token(sqlparse.tokens.Other, f"'\\x{value.hex()}'")
 
@@ -235,7 +235,7 @@ def _parse_placeholder(token):
     if token.value == "%s":
         return Paramstyle.FORMAT, None
 
-    # E.g., %(foo)
+    # E.g., %(foo)s
     matches = re.search(r"%\((\w+)\)s$", token.value)
     if matches:
         return Paramstyle.PYFORMAT, matches.group(1)
@@ -253,6 +253,10 @@ def _escape_verbatim_colons(tokens):
     return tokens
 
 
+def _is_command_token(token):
+    return token.ttype in {sqlparse.tokens.Keyword, sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML}
+
+
 def _is_string_literal(token):
     return token.ttype in [sqlparse.tokens.Literal.String, sqlparse.tokens.Literal.String.Single]
 

src/cs50/sql.py

@@ -14,14 +14,16 @@
 class SQL:
     def __init__(self, url, **engine_kwargs):
         self._session = Session(url, **engine_kwargs)
+        self._dialect = self._session.get_bind().dialect
+        self._is_postgres = self._dialect in {"postgres", "postgresql"}
         self._autocommit = False
 
 
     def execute(self, sql, *args, **kwargs):
         """Execute a SQL statement."""
-        statement = Statement(self._session.get_bind().dialect, sql, *args, **kwargs)
+        statement = Statement(self._dialect, sql, *args, **kwargs)
         command = statement.get_command()
-        if command in ["BEGIN", "START"]:
+        if command in {"BEGIN", "START"}:
             self._autocommit = False
 
         if self._autocommit:
@@ -33,18 +35,15 @@ def execute(self, sql, *args, **kwargs):
             self._session.execute("COMMIT")
             self._session.remove()
 
-        if command in ["COMMIT", "ROLLBACK"]:
+        if command in {"COMMIT", "ROLLBACK"}:
             self._autocommit = True
             self._session.remove()
 
         if command == "SELECT":
             ret = _fetch_select_result(result)
         elif command == "INSERT":
-            if self._session.is_postgres():
-                ret = self._get_last_val()
-            else:
-                ret = result.lastrowid if result.rowcount == 1 else None
-        elif command in ["DELETE", "UPDATE"]:
+            ret = self._last_row_id_or_none(result)
+        elif command in {"DELETE", "UPDATE"}:
             ret = result.rowcount
         else:
             ret = True
@@ -71,6 +70,16 @@ def _execute(self, statement):
             return result
 
 
+    def _last_row_id_or_none(self, result):
+        if self.is_postgres():
+            return self._get_last_val()
+        return result.lastrowid if result.rowcount == 1 else None
+
+
+    def is_postgres(self):
+        return self._is_postgres
+
+
     def _get_last_val(self):
         try:
             return self._session.execute("SELECT LASTVAL()").first()[0]