Improve code organization: docstrings, type hints, and imports

- Add module docstrings to cli.py, admin.py, logging.py, and hash.py
- Add `from __future__ import annotations` for PEP 563 deferred evaluation
- Modernize type hints using Python 3.10+ syntax (X | None instead of Optional[X])
- Use TYPE_CHECKING imports to avoid circular dependencies
- Improve function docstrings to Google style with Args/Returns sections
- Add type annotations to:
  - admin.py: set_password(), kill(), and kill_quick() functions
  - cli.py: cli() function
  - hash.py: all hashing functions
  - logging.py: excepthook() function
  - utils.py: all utility functions and ClassProperty

Co-authored-by: dimitri-yatsenko <dimitri@datajoint.com>

Author: claude

datajoint/admin.py

@@ -1,17 +1,39 @@
+"""
+Administrative utilities for managing database connections and passwords.
+
+This module provides functions for viewing and terminating database connections
+through the MySQL processlist interface, as well as password management.
+"""
+
+from __future__ import annotations
+
 import logging
 from getpass import getpass
 
 import pymysql
 from packaging import version
 
-from .connection import conn
+from .connection import Connection, conn
 from .settings import config
 from .utils import user_choice
 
 logger = logging.getLogger(__name__.split(".")[0])
 
 
-def set_password(new_password=None, connection=None, update_config=None):
+def set_password(
+    new_password: str | None = None,
+    connection: Connection | None = None,
+    update_config: bool | None = None,
+) -> None:
+    """
+    Change the database password for the current user.
+
+    Args:
+        new_password: The new password. If None, prompts for input.
+        connection: A datajoint.Connection object. If None, uses datajoint.conn().
+        update_config: If True, save the new password to local config.
+            If None, prompts the user.
+    """
     connection = conn() if connection is None else connection
     if new_password is None:
         new_password = getpass("New password: ")
@@ -36,20 +58,27 @@ def set_password(new_password=None, connection=None, update_config=None):
         config.save_local(verbose=True)
 
 
-def kill(restriction=None, connection=None, order_by=None):
+def kill(
+    restriction: str | None = None,
+    connection: Connection | None = None,
+    order_by: str | list[str] | None = None,
+) -> None:
     """
-    view and kill database connections.
+    View and interactively kill database connections.
 
-    :param restriction: restriction to be applied to processlist
-    :param connection: a datajoint.Connection object. Default calls datajoint.conn()
-    :param order_by: order by a single attribute or the list of attributes. defaults to 'id'.
+    Displays active database connections matching the optional restriction and
+    prompts the user to select connections to terminate.
 
-    Restrictions are specified as strings and can involve any of the attributes of
-    information_schema.processlist: ID, USER, HOST, DB, COMMAND, TIME, STATE, INFO.
+    Args:
+        restriction: SQL WHERE clause condition to filter the processlist.
+            Can reference any column from information_schema.processlist:
+            ID, USER, HOST, DB, COMMAND, TIME, STATE, INFO.
+        connection: A datajoint.Connection object. If None, uses datajoint.conn().
+        order_by: Column name(s) to sort results by. Defaults to 'id'.
 
     Examples:
-        dj.kill('HOST LIKE "%compute%"') lists only connections from hosts containing "compute".
-        dj.kill('TIME > 600') lists only connections in their current state for more than 10 minutes
+        >>> dj.kill('HOST LIKE "%compute%"')  # connections from hosts containing "compute"
+        >>> dj.kill('TIME > 600')  # connections idle for more than 10 minutes
     """
 
     if connection is None:
@@ -95,18 +124,28 @@ def kill(restriction=None, connection=None, order_by=None):
                     logger.warn("Process not found")
 
 
-def kill_quick(restriction=None, connection=None):
+def kill_quick(
+    restriction: str | None = None,
+    connection: Connection | None = None,
+) -> int:
     """
-    Kill database connections without prompting. Returns number of terminated connections.
+    Kill database connections without prompting.
+
+    Terminates all database connections matching the optional restriction
+    without user confirmation.
 
-    :param restriction: restriction to be applied to processlist
-    :param connection: a datajoint.Connection object. Default calls datajoint.conn()
+    Args:
+        restriction: SQL WHERE clause condition to filter the processlist.
+            Can reference any column from information_schema.processlist:
+            ID, USER, HOST, DB, COMMAND, TIME, STATE, INFO.
+        connection: A datajoint.Connection object. If None, uses datajoint.conn().
 
-    Restrictions are specified as strings and can involve any of the attributes of
-    information_schema.processlist: ID, USER, HOST, DB, COMMAND, TIME, STATE, INFO.
+    Returns:
+        Number of connections terminated.
 
     Examples:
-        dj.kill('HOST LIKE "%compute%"') terminates connections from hosts containing "compute".
+        >>> dj.kill_quick('HOST LIKE "%compute%"')  # kill connections from "compute" hosts
+        >>> dj.kill_quick('TIME > 600')  # kill connections idle for more than 10 minutes
     """
     if connection is None:
         connection = conn()

datajoint/cli.py

@@ -1,16 +1,41 @@
+"""
+Command-line interface for DataJoint Python.
+
+This module provides a console interface for interacting with DataJoint databases,
+allowing users to connect to servers and work with virtual modules from the command line.
+
+Usage:
+    datajoint [-u USER] [-p PASSWORD] [-h HOST] [-s SCHEMA:MODULE ...]
+
+Example:
+    datajoint -u root -h localhost -s mydb:experiment mydb:subject
+"""
+
+from __future__ import annotations
+
 import argparse
 from code import interact
 from collections import ChainMap
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
 
 import datajoint as dj
 
 
-def cli(args: list = None):
+def cli(args: Sequence[str] | None = None) -> None:
     """
-    Console interface for DataJoint Python
+    Console interface for DataJoint Python.
+
+    Launches an interactive Python shell with DataJoint configured and optional
+    virtual modules loaded for database schemas.
+
+    Args:
+        args: List of command-line arguments. If None, reads from sys.argv.
 
-    :param args: List of arguments to be passed in, defaults to reading stdin
-    :type args: list, optional
+    Raises:
+        SystemExit: Always raised when the interactive session ends.
     """
     parser = argparse.ArgumentParser(
         prog="datajoint",

datajoint/hash.py

@@ -1,26 +1,56 @@
+"""
+Hashing utilities for DataJoint.
+
+This module provides functions for computing hashes of data, streams, and files.
+These are used for checksums, content-addressable storage, and primary key hashing.
+"""
+
+from __future__ import annotations
+
 import hashlib
 import io
 import uuid
+from collections.abc import Mapping
 from pathlib import Path
+from typing import IO
 
 
-def key_hash(mapping):
+def key_hash(mapping: Mapping) -> str:
     """
-    32-byte hash of the mapping's key values sorted by the key name.
-    This is often used to convert a long primary key value into a shorter hash.
-    For example, the JobTable in datajoint.jobs uses this function to hash the primary key of autopopulated tables.
+    Compute a 32-character hex hash of a mapping's values, sorted by key name.
+
+    This is commonly used to convert long primary key values into shorter hashes.
+    For example, the JobTable in datajoint.jobs uses this function to hash the
+    primary keys of autopopulated tables.
+
+    Args:
+        mapping: A dict-like object whose values will be hashed.
+
+    Returns:
+        A 32-character hexadecimal MD5 hash string.
+
+    Example:
+        >>> key_hash({'subject_id': 1, 'session': 5})
+        'a1b2c3d4e5f6...'
     """
     hashed = hashlib.md5()
     for k, v in sorted(mapping.items()):
         hashed.update(str(v).encode())
     return hashed.hexdigest()
 
 
-def uuid_from_stream(stream, *, init_string=""):
+def uuid_from_stream(stream: IO[bytes], *, init_string: str = "") -> uuid.UUID:
     """
-    :return: 16-byte digest of stream data
-    :stream: stream object or open file handle
-    :init_string: string to initialize the checksum
+    Compute a UUID from the contents of a binary stream.
+
+    Reads the stream in chunks and computes an MD5 hash, returning it as a UUID.
+
+    Args:
+        stream: A binary stream object (file handle opened in 'rb' mode or BytesIO).
+        init_string: Optional string to initialize the hash (acts as a salt).
+
+    Returns:
+        A UUID object derived from the MD5 hash of the stream contents.
     """
     hashed = hashlib.md5(init_string.encode())
     chunk = True
@@ -31,9 +61,29 @@ def uuid_from_stream(stream, *, init_string=""):
     return uuid.UUID(bytes=hashed.digest())
 
 
-def uuid_from_buffer(buffer=b"", *, init_string=""):
+def uuid_from_buffer(buffer: bytes = b"", *, init_string: str = "") -> uuid.UUID:
+    """
+    Compute a UUID from a bytes buffer.
+
+    Args:
+        buffer: The binary data to hash.
+        init_string: Optional string to initialize the hash (acts as a salt).
+
+    Returns:
+        A UUID object derived from the MD5 hash of the buffer.
+    """
     return uuid_from_stream(io.BytesIO(buffer), init_string=init_string)
 
 
-def uuid_from_file(filepath, *, init_string=""):
+def uuid_from_file(filepath: str | Path, *, init_string: str = "") -> uuid.UUID:
+    """
+    Compute a UUID from the contents of a file.
+
+    Args:
+        filepath: Path to the file to hash.
+        init_string: Optional string to initialize the hash (acts as a salt).
+
+    Returns:
+        A UUID object derived from the MD5 hash of the file contents.
+    """
     return uuid_from_stream(Path(filepath).open("rb"), init_string=init_string)

datajoint/logging.py

@@ -1,6 +1,18 @@
+"""
+Logging configuration for the DataJoint package.
+
+This module sets up the default logging handler and format for DataJoint,
+and provides a custom exception hook to log uncaught exceptions.
+
+The log level can be configured via the DJ_LOG_LEVEL environment variable.
+"""
+
+from __future__ import annotations
+
 import logging
 import os
 import sys
+from types import TracebackType
 
 logger = logging.getLogger(__name__.split(".")[0])
 
@@ -15,7 +27,22 @@
 logger.handlers = [stream_handler]
 
 
-def excepthook(exc_type, exc_value, exc_traceback):
+def excepthook(
+    exc_type: type[BaseException],
+    exc_value: BaseException,
+    exc_traceback: TracebackType | None,
+) -> None:
+    """
+    Custom exception hook that logs uncaught exceptions.
+
+    Keyboard interrupts are passed to the default handler; all other exceptions
+    are logged as errors with full traceback information.
+
+    Args:
+        exc_type: The exception class.
+        exc_value: The exception instance.
+        exc_traceback: The traceback object.
+    """
     if issubclass(exc_type, KeyboardInterrupt):
         sys.__excepthook__(exc_type, exc_value, exc_traceback)
         return