complete icon support

Author: pvsr

qbpm.1.scd

@@ -77,6 +77,18 @@ $HOME/.local/share), but this can be set to another directory by passing
 		qbpm launch -n qb-dev --debug --json-logging
 		```
 
+*icon* [options] <profile> <icon>
+	Install an icon to _profile_. _icon_ may be a url, a path to an image file,
+	or, if --by-name is passed, the name of an xdg icon installed on your
+	system. If _icon_ is a url, qbpm will fetch the page and attempt to find a
+	suitable favicon.
+
+	Options:
+
+	*-n, --by-name*
+		Interpret _icon_ as the name of an xdg icon file according to the
+		freedesktop.org icon specification. Likely to only work on Linux.
+
 *choose* [options]
 	Open a menu to choose a qutebrowser profile to launch. On linux this defaults
 	to dmenu or another compatible menu program such as rofi, and on macOS this

qbpm/config.py

@@ -0,0 +1,3 @@
+# TODO real config file
+default_icon = "qutebrowser"
+application_name_suffix = " (qutebrowser profile)"

qbpm/icons.py

@@ -1,5 +1,7 @@
+import re
 import shutil
 from collections import namedtuple
+from collections.abc import Iterator
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Optional
@@ -9,56 +11,118 @@
 import requests
 from PIL import Image
 
-from .profiles import Profile
+from . import Profile, __version__
+from .utils import debug, error
 
 PREFERRED_ICONS = [
-    "favicon.ico",
-    "favicon.svg",
-    "favicon.png",
-    ".ico",
+    re.compile(p)
+    for p in [
+        r"favicon.*\.ico$",
+        r"favicon.*\.svg$",
+        r"favicon.*\.png$",
+        r"\.ico$",
+        r"icon\.png$",
+        r"\.svg$",
+    ]
 ]
 
+# https://github.com/scottwernervt/favicon/blob/123e431f53b2c4903b540246a85db0b1633d4786/src/favicon/favicon.py#L40
 Icon = namedtuple("Icon", ["url", "width", "height", "format"])
 
 
 def choose_icon(icons: list[Icon]) -> Optional[Icon]:
-    print(f"{icons=}")
-    for key in PREFERRED_ICONS:
+    for pattern in PREFERRED_ICONS:
         for icon in icons:
-            if key in icon.url:
-                print(f"chose {icon}")
+            if pattern.search(icon.url):
+                debug(f"chose {icon.url}")
                 return icon
     return None
 
 
-def download_icon(profile: Profile, home_page: str) -> Optional[Path]:
+headers = {"user-agent": f"qbpm/{__version__}"}
+
+
+def download_icon(profile: Profile, home_page: str, overwrite: bool) -> Optional[Path]:
     if not urlparse(home_page).scheme:
         home_page = f"https://{home_page}"
-    icons = favicon.get(home_page, timeout=10)
+    try:
+        icons = favicon.get(home_page, headers=headers, timeout=10)
+    except Exception as e:
+        debug(str(e))
+        error(f"failed to fetch favicon from {home_page}")
+        return None
+    if not icons:
+        error(f"no favicon found on {home_page}")
+        return None
     icon = choose_icon(icons)
     if not icon:
-        # TODO print
+        print(f"no favicons found matching one of {PREFERRED_ICONS}")
         return None
 
     tmp_dir = TemporaryDirectory()
     work_dir = Path(tmp_dir.name)
     work_icon = work_dir / f"favicon.{icon.format}"
-    resp = requests.get(icon.url, timeout=10)
+    icon_body = requests.get(icon.url, headers=headers, timeout=10)
     with work_icon.open("wb") as icon_file:
-        for chunk in resp.iter_content(1024):
+        for chunk in icon_body.iter_content(1024):
             icon_file.write(chunk)
 
-    icon_path = profile.root / f"icon.{icon.format}"
-    if icon.format == "ico":
-        icon_path = icon_path.with_suffix(".png")
-        image = Image.open(work_icon)
-        image.save(icon_path)
+    return install_icon_file(profile, work_icon, overwrite, icon.url)
+
+
+def icon_for_profile(profile: Profile) -> Optional[str]:
+    icon_file = next(find_icon_files(profile), None)
+    if icon_file and icon_file.suffix == ".name":
+        return icon_file.read_text()
+    return str(icon_file) if icon_file else None
+
+
+def install_icon_file(
+    profile: Profile, src: Path, overwrite: bool, origin: Optional[str] = None
+) -> Optional[Path]:
+    if not check_or_replace_existing_icon(profile, overwrite):
+        return None
+    icon_format = src.suffix
+    dest = profile.root / f"icon{icon_format}"
+    if icon_format == ".ico":
+        dest = dest.with_suffix(".png")
+        try:
+            image = Image.open(src)
+            image.save(dest, icon_format="png")
+        except Exception as e:
+            debug(str(e))
+            error(f"failed to convert {origin or src} to png")
+            return None
     else:
-        shutil.move(work_icon, icon_path)
+        shutil.copy(src, dest)
+    return dest.absolute()
+
 
-    return icon_path.absolute()
+def install_icon_by_name(profile: Profile, icon_name: str, overwrite: bool) -> bool:
+    if not check_or_replace_existing_icon(profile, overwrite):
+        return False
+    file = profile.root / "icon.name"
+    file.write_text(icon_name)
+    return True
+
+
+def check_or_replace_existing_icon(profile: Profile, overwrite: bool) -> bool:
+    existing_icons = find_icon_files(profile)
+    existing_icon = next(existing_icons, None)
+    if not existing_icon:
+        return True
+    elif overwrite:
+        existing_icon.unlink()
+        for icon in existing_icons:
+            icon.unlink()
+        return True
+    else:
+        error(f"icon already exists in {profile.root}, pass --overwrite to replace it")
+        return False
 
 
-def find_icon_file(profile: Profile) -> Optional[Path]:
-    icon_path = next(filter(lambda p: p.is_file(), profile.root.glob("icon.*")), None)
-    return icon_path.absolute() if icon_path else icon_path
+def find_icon_files(profile: Profile) -> Iterator[Path]:
+    for ext in ["png", "svg", "name"]:
+        icon = profile.root / f"icon.{ext}"
+        if icon.is_file():
+            yield icon.absolute()

qbpm/main.py

@@ -113,14 +113,20 @@ def launch(profile_dir: Path, profile_name: str, **kwargs: Any) -> None:
 @click.option(
     "-f", "--foreground", is_flag=True, help="Run qutebrowser in the foreground."
 )
+@click.option(
+    "-i",
+    "--icon",
+    "force_icon",
+    is_flag=True,
+    help="Attach icons to menu items using rofi's extended dmenu spec even if your menu is not known to support it. Only works if at least one profile has an icon installed.",
+)
 @click.pass_obj
 def choose(profile_dir: Path, **kwargs: Any) -> None:
     """Choose a profile to launch.
     Support is built in for many X and Wayland launchers, as well as applescript dialogs.
     All QB_ARGS are passed on to qutebrowser."
     """
-    # TODO force-icon
-    exit_with(operations.choose(profile_dir=profile_dir, force_icon=False, **kwargs))
+    exit_with(operations.choose(profile_dir=profile_dir, **kwargs))
 
 
 @main.command()
@@ -157,26 +163,21 @@ def desktop(
 
 @main.command()
 @click.argument("profile_name")
-@click.argument("icon")
+@click.argument("icon", metavar="ICON_LOCATION")
+@click.option(
+    "-n",
+    "--by-name",
+    is_flag=True,
+    help="interpret ICON_LOCATION as the name of an icon in an installed icon theme instead of a file or url.",
+)
+@click.option("--overwrite", is_flag=True, help="Replace the current icon.")
 @click.pass_obj
-def icon(profile_dir: Path, profile_name: str, icon: str) -> None:
-    # TODO support --all?
-    """Install an icon for the profile."""
+def icon(profile_dir: Path, profile_name: str, **kwargs: Any) -> None:
+    """Install an icon for the profile. ICON_LOCATION may be a url to download a favicon from,
+    a path to an image file, or, with --by-name, the name of an xdg icon installed on your system.
+    """
     profile = Profile(profile_name, profile_dir)
-    if not profile.exists():
-        error(f"profile {profile.name} not found at {profile.root}")
-        sys.exit(1)
-    icon_file: Optional[Path]
-    if Path(icon).is_file():
-        # TODO move to profile dir
-        # TODO overwrite protection
-        icon_file = Path(icon)
-    else:
-        icon_file = profiles.download_icon(profile, icon)
-    # TODO check if desktop file exists and should be updated
-    desktop_file = False
-    if desktop_file:
-        profiles.create_desktop_file(profile, icon_file)
+    exit_with(operations.icon(profile, **kwargs))
 
 
 def then_launch(

qbpm/operations.py

@@ -7,8 +7,7 @@
 
 from xdg import BaseDirectory
 
-from . import Profile, profiles
-from .icons import find_icon_file
+from . import Profile, config, icons, profiles
 from .utils import AUTO_MENUS, error, installed_menus
 
 
@@ -59,12 +58,29 @@ def launch(
 def desktop(profile: Profile) -> bool:
     exists = profile.exists()
     if exists:
-        profiles.create_desktop_file(profile)
+        profiles.create_desktop_file(profile, icons.icon_for_profile(profile))
     else:
         error(f"profile {profile.name} not found at {profile.root}")
     return exists
 
 
+def icon(profile: Profile, icon: str, by_name: bool, overwrite: bool) -> bool:
+    if not profile.exists():
+        error(f"profile {profile.name} not found at {profile.root}")
+        return False
+    if by_name:
+        icon_id = icon if icons.install_icon_by_name(profile, icon, overwrite) else None
+    else:
+        if Path(icon).is_file():
+            icon_file = icons.install_icon_file(profile, Path(icon), overwrite)
+        else:
+            icon_file = icons.download_icon(profile, icon, overwrite)
+        icon_id = str(icon_file) if icon_file else None
+    if icon_id:
+        profiles.add_to_desktop_file(profile, "Icon", icon_id)
+    return icon_id is not None
+
+
 def choose(
     profile_dir: Path,
     menu: str,
@@ -158,10 +174,10 @@ def menu_command(
 
 
 def build_menu_items(profiles: list[Profile], icon: bool) -> str:
-    if icon and any(icons := [find_icon_file(p) for p in profiles]):
+    if icon and any(profile_icons := [icons.icon_for_profile(p) for p in profiles]):
         menu_items = [
-            f"{p.name}\0icon\x1f{icon or 'qutebrowser'}"
-            for (p, icon) in zip(profiles, icons)
+            f"{p.name}\0icon\x1f{icon or config.default_icon}"
+            for (p, icon) in zip(profiles, profile_icons)
         ]
     else:
         menu_items = [p.name for p in profiles]

qbpm/profiles.py

@@ -5,8 +5,7 @@
 from xdg import BaseDirectory
 from xdg.DesktopEntry import DesktopEntry
 
-from . import Profile
-from .icons import download_icon
+from . import Profile, config, icons
 from .utils import error, user_config_dir
 
 
@@ -44,18 +43,26 @@ def create_config(
 application_dir = Path(BaseDirectory.xdg_data_home) / "applications" / "qbpm"
 
 
-def create_desktop_file(profile: Profile, icon: Path | str | None = None) -> None:
+def create_desktop_file(profile: Profile, icon: Optional[str] = None) -> None:
     desktop = DesktopEntry(str(application_dir / f"{profile.name}.desktop"))
-    desktop.set("Name", f"{profile.name} (qutebrowser profile)")
-    # TODO allow passing in an icon value
-    desktop.set("Icon", icon or "qutebrowser")
+    desktop.set("Name", f"{profile.name}{config.application_name_suffix}")
+    desktop.set("Icon", icon or config.default_icon)
     desktop.set("Exec", " ".join(profile.cmdline()) + " %u")
     desktop.set("Categories", ["Network"])
     desktop.set("Terminal", False)
     desktop.set("StartupNotify", True)
     desktop.write()
 
 
+def add_to_desktop_file(profile: Profile, key: str, value: str) -> None:
+    desktop_file = application_dir / f"{profile.name}.desktop"
+    if not desktop_file.exists():
+        return
+    desktop = DesktopEntry(str(application_dir / f"{profile.name}.desktop"))
+    desktop.set(key, value)
+    desktop.write()
+
+
 def ensure_profile_exists(profile: Profile, create: bool = True) -> bool:
     if profile.root.exists() and not profile.root.is_dir():
         error(f"{profile.root} is not a directory")
@@ -78,8 +85,8 @@ def new_profile(
         create_config(profile, home_page, overwrite)
         if home_page:
             # TODO catch errors?
-            icon = download_icon(profile, home_page)
+            icon = icons.download_icon(profile, home_page, overwrite)
         if desktop_file:
-            create_desktop_file(profile, icon)
+            create_desktop_file(profile, str(icon) if icon else None)
         return True
     return False

qbpm/utils.py

@@ -12,6 +12,15 @@
 AUTO_MENUS = WAYLAND_MENUS + X11_MENUS
 SUPPORTED_MENUS = [*AUTO_MENUS, "fzf", "applescript"]
 
+global debugging
+debugging: bool = True
+
+
+# TODO replace with actual logger
+def debug(msg: str) -> None:
+    if debugging:
+        print(f"debug: {msg}", file=stderr)
+
 
 def error(msg: str) -> None:
     print(f"error: {msg}", file=stderr)