Add basic IPython config and startup scripts

9 files changed, 472 insertions, 0 deletions

diff --git a/dot_ipython/profile_default/ipython_config.py b/dot_ipython/profile_default/ipython_config.py
new file mode 100644
index 0000000..07ab804
--- /dev/null
+++ b/dot_ipython/profile_default/ipython_config.py
@@ -0,0 +1,59 @@
+"""IPython config."""
+
+import datetime
+import importlib
+import sys
+
+import IPython
+from IPython.core.interactiveshell import InteractiveShell
+from IPython.core.shellapp import InteractiveShellApp
+from IPython.terminal.prompts import Prompts
+from pygments.token import Token
+from traitlets.config import Config
+
+c: Config
+c.InteractiveShell: InteractiveShell
+c.InteractiveShellApp: InteractiveShellApp
+
+
+class TimePrompt(Prompts):
+    def in_prompt_tokens(self) -> list[tuple[Token, str]]:
+        return [
+            (Token.Prompt, datetime.datetime.now().strftime("%H:%M")),
+            (Token.Prompt, " ["),
+            (Token.PromptNum, str(self.shell.execution_count)),
+            (Token.Prompt, "]: "),
+        ]
+
+    def out_prompt_tokens(self) -> list[tuple[Token, str]]:
+        return [
+            (Token.OutPrompt, datetime.datetime.now().strftime("%H:%M")),
+            (Token.OutPrompt, " ["),
+            (Token.OutPromptNum, str(self.shell.execution_count)),
+            (Token.OutPrompt, "]: "),
+        ]
+
+
+version_parts = [f"IPython {IPython.__version__}"]
+for key_library, name in (
+    ("pandas", "pd"),
+    ("numpy", "np"),
+    ("numba", "nb"),
+    ("sqlalchemy", "sa"),
+    ("pydantic", "pydantic"),
+):
+    try:
+        module = importlib.import_module(key_library)
+    except ModuleNotFoundError:
+        pass
+    else:
+        version_parts.append(f"{name} {module.__version__}")
+versions = ", ".join(version_parts)
+
+c.InteractiveShellApp.exec_lines = ["%autoreload 2"]
+c.InteractiveShellApp.extensions = ["autoreload"]
+c.InteractiveShell.history_length = 100_000
+c.InteractiveShell.history_load_length = 100_000
+c.InteractiveShell.prompts_class = TimePrompt
+c.InteractiveShell.banner1 = f"Python {sys.version.split()[0]} [{versions}]\n"
+c.InteractiveShell.banner2 = ""
diff --git a/dot_ipython/profile_default/startup/00_internal.py b/dot_ipython/profile_default/startup/00_internal.py
new file mode 100644
index 0000000..57403d8
--- /dev/null
+++ b/dot_ipython/profile_default/startup/00_internal.py
@@ -0,0 +1,32 @@
+"""Internal startup utils."""
+
+_IPYTHON_INTERNAL_NAMES = set(locals())
+_INTERNAL_NAMES = set(locals())
+
+
+def _summarise_startup() -> None:
+    import inspect
+    import logging
+    import os
+
+    frame = inspect.stack()[1].frame
+    filename = os.path.basename(frame.f_locals["__file__"])
+    nice_name, _ = os.path.splitext(filename)
+    nice_name = nice_name.lstrip("0123456789_")
+
+    names = sorted(
+        name
+        for name in set(frame.f_locals) - _INTERNAL_NAMES
+        if not name.startswith("_")
+    )
+
+    if names:
+        for name in names:
+            _INTERNAL_NAMES.add(name)
+        logging.getLogger("startup").setLevel(logging.INFO)
+        logger = logging.getLogger(f"startup.{nice_name}")
+        log_level = getattr(logging, os.getenv("STARTUP_LOG_LEVEL", "DEBUG"))
+        logger.log(log_level, "%s", ", ".join(names))
+
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/10_colours.py b/dot_ipython/profile_default/startup/10_colours.py
new file mode 100644
index 0000000..eb6e5a8
--- /dev/null
+++ b/dot_ipython/profile_default/startup/10_colours.py
@@ -0,0 +1,56 @@
+"""Simple colours for colourful colouring."""
+
+
+class FG:
+    RESET = "\u001b[0m"
+
+    DARK_BLACK = BLACK = "\u001b[30m"
+    DARK_RED = "\u001b[31m"
+    DARK_GREEN = "\u001b[32m"
+    DARK_YELLOW = BROWN = "\u001b[33m"
+    DARK_BLUE = BLUE = "\u001b[34m"
+    DARK_MAGENTA = PURPLE = "\u001b[35m"
+    DARK_CYAN = "\u001b[36m"
+    DARK_WHITE = "\u001b[37m"
+
+    BRIGHT_BLACK = GRAY = GREY = "\u001b[30;1m"
+    BRIGHT_RED = RED = "\u001b[31;1m"
+    BRIGHT_GREEN = GREEN = "\u001b[32;1m"
+    BRIGHT_YELLOW = YELLOW = "\u001b[33;1m"
+    BRIGHT_BLUE = "\u001b[34;1m"
+    BRIGHT_MAGENTA = MAGENTA = "\u001b[35;1m"
+    BRIGHT_CYAN = CYAN = "\u001b[36;1m"
+    BRIGHT_WHITE = WHITE = "\u001b[37;1m"
+
+    def __getitem__(self, rgb):
+        r, g, b = rgb
+        return f"\u001b[38;2;{r};{g};{b}m"
+
+
+class BG:
+    RESET = "\u001b[0m"
+
+    DARK_BLACK = BLACK = "\u001b[40m"
+    DARK_RED = "\u001b[41m"
+    DARK_GREEN = "\u001b[42m"
+    DARK_YELLOW = BROWN = "\u001b[44m"
+    DARK_BLUE = BLUE = "\u001b[44m"
+    DARK_MAGENTA = PURPLE = "\u001b[45m"
+    DARK_CYAN = "\u001b[46m"
+    DARK_WHITE = "\u001b[47m"
+
+    BRIGHT_BLACK = GRAY = GREY = "\u001b[40;1m"
+    BRIGHT_RED = RED = "\u001b[41;1m"
+    BRIGHT_GREEN = GREEN = "\u001b[42;1m"
+    BRIGHT_YELLOW = YELLOW = "\u001b[44;1m"
+    BRIGHT_BLUE = "\u001b[44;1m"
+    BRIGHT_MAGENTA = MAGENTA = "\u001b[45;1m"
+    BRIGHT_CYAN = CYAN = "\u001b[46;1m"
+    BRIGHT_WHITE = WHITE = "\u001b[47;1m"
+
+    def __getitem__(self, rgb):
+        r, g, b = rgb
+        return f"\u001b[48;2;{r};{g};{b}m"
+
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/20_logging.py b/dot_ipython/profile_default/startup/20_logging.py
new file mode 100644
index 0000000..86ef226
--- /dev/null
+++ b/dot_ipython/profile_default/startup/20_logging.py
@@ -0,0 +1,111 @@
+"""Set some (hopefully) sane and comfortable logging defaults."""
+
+import logging
+
+
+def _setup() -> logging.Logger:
+    import logging
+    import os
+    import setuptools
+    import glob
+
+    # Determine log levels
+    project_level = getattr(logging, os.environ.get("LOG_LEVEL", "INFO"))
+    root_level = getattr(logging, os.environ.get("ROOT_LOG_LEVEL", "WARNING"))
+    root_level = max(project_level, root_level)
+
+    # Determine projects
+    project_names = (
+        ["__main__", "repl"]
+        + [x[:-3] for x in glob.glob("*.py")]
+        + setuptools.find_packages()
+    )
+
+    # Set logging levels
+    logging.root.setLevel(root_level)
+    for name in project_names:
+        logging.getLogger(name).setLevel(project_level)
+
+    # Configure logging format
+    try:
+        from colorlog import ColoredFormatter
+    except ModuleNotFoundError:
+        formatter = logging.Formatter(
+            "\u001b[36m%(asctime)s "
+            "\u001b[0m%(levelname)-8s "
+            "\u001b[32m%(threadName)s "
+            "\u001b[36m%(name)s "
+            "\u001b[37;1m%(message)s"
+            "\u001b[0m",
+            datefmt="%H:%M",
+        )
+    else:
+
+        class FancyFormatter(logging.Formatter):
+            """Formatter that uses a different formatter for debug logs."""
+
+            def __init__(
+                self,
+                formatter: logging.Formatter,
+                debug_formatter: logging.Formatter,
+            ) -> None:
+                super().__init__()
+                self.formatter = formatter
+                self.debug_formatter = debug_formatter
+
+            def format(
+                self,
+                record: logging.LogRecord,
+            ) -> str:
+                if record.levelno <= logging.DEBUG:
+                    return self.debug_formatter.format(record)
+                return self.formatter.format(record)
+
+        fmt = (
+            "%(cyan)s%(asctime)s "
+            "%(log_color)s%(levelname)-8s "
+            "%(green)s%(threadName)s "
+            "%(cyan)s%(name)s "
+            "%(white)s%(message)s"
+        )
+        formatter = FancyFormatter(
+            formatter=ColoredFormatter(
+                fmt=fmt,
+                datefmt="%H:%M",
+                reset=True,
+                log_colors={
+                    "DEBUG": "white",
+                    "INFO": "white",
+                    "WARNING": "yellow",
+                    "ERROR": "red",
+                    "CRITICAL": "bold_red",
+                },
+            ),
+            debug_formatter=ColoredFormatter(
+                fmt=f"%(thin)s{fmt}",
+                datefmt="%H:%M",
+                reset=True,
+                log_colors={
+                    "DEBUG": "white",
+                    "INFO": "white",
+                    "WARNING": "yellow",
+                    "ERROR": "red",
+                    "CRITICAL": "bold_red",
+                },
+            ),
+        )
+
+    # Create root handler with formatter
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    if not logging.root.handlers:
+        logging.root.addHandler(handler)
+
+    return logging.getLogger("repl")
+
+
+logger = _setup()
+
+del _setup
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/40_json_literals.py b/dot_ipython/profile_default/startup/40_json_literals.py
new file mode 100644
index 0000000..f85d3d5
--- /dev/null
+++ b/dot_ipython/profile_default/startup/40_json_literals.py
@@ -0,0 +1,7 @@
+"""Make pasting JSON easier by providing convenient literals."""
+
+true = True
+false = False
+null = None
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/45_datetime.py b/dot_ipython/profile_default/startup/45_datetime.py
new file mode 100644
index 0000000..d8dcf4e
--- /dev/null
+++ b/dot_ipython/profile_default/startup/45_datetime.py
@@ -0,0 +1,99 @@
+"""Datetime utils."""
+
+import datetime
+import zoneinfo
+
+# Add some pytz timezones, if pytz available
+try:
+    import pytz
+except ModuleNotFoundError:
+    try:
+        UTC = zoneinfo.ZoneInfo("UTC")
+        LDN = zoneinfo.ZoneInfo("Europe/London")
+        NYC = zoneinfo.ZoneInfo("America/New_York")
+        GST = zoneinfo.ZoneInfo("Asia/Dubai")
+    except KeyError:
+        UTC = datetime.timezone.utc
+        GST = datetime.timezone(datetime.timedelta(hours=4), "Asia/Dubai")
+else:
+    UTC = pytz.timezone("UTC")
+    LDN = pytz.timezone("Europe/London")
+    NYC = pytz.timezone("America/New_York")
+    GST = pytz.timezone("Asia/Dubai")
+
+
+class T(datetime.date):
+    """Convenient dev date wrapper."""
+
+    def __add__(self, other):
+        if other in (D, H, M, S, B):
+            other = other(1)
+        if isinstance(other, int):
+            other = datetime.timedelta(days=other)
+        v = super().__add__(other)
+        if isinstance(v, datetime.date) and not isinstance(
+            v, datetime.datetime
+        ):
+            return type(self)(v.year, v.month, v.day)
+        return v
+
+    def __sub__(self, other):
+        if other in (D, H, M, S, B):
+            other = other(1)
+        if isinstance(other, int):
+            other = datetime.timedelta(days=other)
+        v = super().__sub__(other)
+        if isinstance(v, datetime.date) and not isinstance(
+            v, datetime.datetime
+        ):
+            return type(self)(v.year, v.month, v.day)
+        return v
+
+    def __call__(self, *args):
+        return type(self)(*args)
+
+
+_T = datetime.date.today()
+T = T(_T.year, _T.month, _T.day)
+D = lambda d: datetime.timedelta(days=d)
+H = lambda h: datetime.timedelta(hours=h)
+M = lambda m: datetime.timedelta(minutes=m)
+S = lambda s: datetime.timedelta(seconds=s)
+B = lambda s: BDelta(s)
+
+
+def now(tz=UTC):
+    return datetime.datetime.now(tz=tz)
+
+
+class BDelta:
+    """Convenient dev weekday delta."""
+
+    def __init__(self, n: int):
+        self.n = n
+
+    def __add__(self, other):
+        n = self.n
+        if isinstance(other, datetime.date):
+            if n >= 0:
+                op = lambda a, b: a + b
+            else:
+                op = lambda a, b: a - b
+            while n:
+                while other.weekday() in (5, 6):
+                    other = op(other, datetime.timedelta(days=1))
+                other = op(other, datetime.timedelta(days=1))
+                n = op(n, -1)
+            while other.weekday() in (5, 6):
+                other = op(other, datetime.timedelta(days=1))
+            return T(other.year, other.month, other.day)
+        return NotImplemented
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __rsub__(self, other):
+        return other + B(-self.n)
+
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/80_pandas.py b/dot_ipython/profile_default/startup/80_pandas.py
new file mode 100644
index 0000000..be18869
--- /dev/null
+++ b/dot_ipython/profile_default/startup/80_pandas.py
@@ -0,0 +1,70 @@
+"""Pandas utilities."""
+
+
+def _setup():
+    import logging
+    import time
+    import fcntl
+    import struct
+    import termios
+    import threading
+
+    class PandasResizer(threading.Thread):
+        """Pandas width/max_row resizer depending on tty size."""
+
+        def __init__(self, max_rows=200):
+            self.max_rows = max_rows
+            self.width = None
+            self.logger = logging.getLogger("startup.pandas")
+            self._run = False
+            super().__init__(name=type(self).__name__, daemon=True)
+
+        @staticmethod
+        def _size():
+            ioctl = fcntl.ioctl(
+                0, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)
+            )
+            th, tw, hp, wp = struct.unpack("HHHH", ioctl)
+            return tw, th
+
+        def run(self):
+            last_resize = float("inf")
+            first_print = True
+            self.logger.info("Resizer running")
+            self._run = True
+            while self._run:
+                width = self._size()[0] - 1
+                if width != self.width:
+                    last_resize = time.time()
+                    self.width = width
+                if first_print or last_resize + 0.5 < time.time():
+                    last_resize = float("inf")
+                    pd.set_option("display.width", self.width)
+                    pd.set_option("display.max_rows", self.max_rows)
+                    if first_print:
+                        self.logger.info(
+                            "Resized pandas to %dx%d",
+                            self.width,
+                            self.max_rows,
+                        )
+                        first_print = False
+                time.sleep(0.2)
+
+        def stop(self):
+            self._run = False
+
+    PandasResizer().start()
+
+
+try:
+    import pandas as pd
+    import numpy as np
+except ModuleNotFoundError:
+    pass
+else:
+    _setup()
+
+
+del _setup
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/95_imports.py b/dot_ipython/profile_default/startup/95_imports.py
new file mode 100644
index 0000000..a4857ac
--- /dev/null
+++ b/dot_ipython/profile_default/startup/95_imports.py
@@ -0,0 +1,32 @@
+"""Some useful imports."""
+
+
+import sys
+import pickle
+import json
+from collections import (
+    Counter,
+    defaultdict,
+)
+from itertools import (
+    count,
+    product,
+    starmap,
+)
+from pathlib import Path
+
+try:
+    from devtools import debug as _debug
+
+    def dprint(*args, **kwargs):
+        _debug(*args, **kwargs)
+
+except ModuleNotFoundError:
+    pass
+
+try:
+    import yaml
+except ModuleNotFoundError:
+    pass
+
+_summarise_startup()
diff --git a/dot_ipython/profile_default/startup/99_internal.py b/dot_ipython/profile_default/startup/99_internal.py
new file mode 100644
index 0000000..2a594eb
--- /dev/null
+++ b/dot_ipython/profile_default/startup/99_internal.py
@@ -0,0 +1,6 @@
+"""Internal startup cleanup."""
+
+del _IPYTHON_INTERNAL_NAMES
+del _INTERNAL_NAMES
+
+del _summarise_startup