wip

Author: pvsr

flake.nix

@@ -21,13 +21,15 @@
             // {
               buildInputs = with pkgs; [
                 ruff
-                xdg-utils.out
+                xdg-utils
                 (python3.withPackages (ps:
                   with ps; [
                     pyxdg
                     favicon
                     pillow
+                    types-pillow
                     requests
+                    types-requests
                     click
 
                     pytest

pyproject.toml

@@ -42,7 +42,7 @@ warn_return_any = true
 warn_unused_ignores = true
 
 [[tool.mypy.overrides]]
-module = "xdg.*"
+module = ["xdg.*", "favicon.*"]
 ignore_missing_imports = true
 
 [[tool.mypy.overrides]]

qbpm/icons.py

@@ -0,0 +1,64 @@
+import shutil
+from collections import namedtuple
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Optional
+from urllib.parse import urlparse
+
+import favicon
+import requests
+from PIL import Image
+
+from .profiles import Profile
+
+PREFERRED_ICONS = [
+    "favicon.ico",
+    "favicon.svg",
+    "favicon.png",
+    ".ico",
+]
+
+Icon = namedtuple("Icon", ["url", "width", "height", "format"])
+
+
+def choose_icon(icons: list[Icon]) -> Optional[Icon]:
+    print(f"{icons=}")
+    for key in PREFERRED_ICONS:
+        for icon in icons:
+            if key in icon.url:
+                print(f"chose {icon}")
+                return icon
+    return None
+
+
+def download_icon(profile: Profile, home_page: str) -> Optional[Path]:
+    if not urlparse(home_page).scheme:
+        home_page = f"https://{home_page}"
+    icons = favicon.get(home_page, timeout=10)
+    icon = choose_icon(icons)
+    if not icon:
+        # TODO print
+        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)
+    with work_icon.open("wb") as icon_file:
+        for chunk in resp.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)
+    else:
+        shutil.move(work_icon, icon_path)
+
+    return icon_path.absolute()
+
+
+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

qbpm/main.py

@@ -20,6 +20,7 @@ def creator_options(f: Callable[..., Any]) -> Callable[..., Any]:
                 is_flag=True,
                 help="If --launch is set, run qutebrowser in the foreground.",
             ),
+            # TODO --no-icon?
             click.option(
                 "--no-desktop-file",
                 "desktop_file",
@@ -130,7 +131,8 @@ def choose(profile_dir: Path, **kwargs: Any) -> None:
     Support is built in for many X and Wayland launchers, as well as applescript dialogs.
     All QB_ARGS are passed on to qutebrowser."
     """
-    exit_with(operations.choose(profile_dir=profile_dir, **kwargs))
+    # TODO force-icon
+    exit_with(operations.choose(profile_dir=profile_dir, force_icon=False, **kwargs))
 
 
 @main.command()
@@ -155,14 +157,26 @@ def list_(profile_dir: Path) -> None:
 
 @main.command()
 @click.argument("profile_name")
+@click.argument("icon")
 @click.pass_obj
-def icon(profile_dir: Path, profile_name: str) -> None:
-    """Edit a profile's config.py."""
+def icon(profile_dir: Path, profile_name: str, icon: str) -> None:
+    # TODO support --all?
+    """Install an icon for the profile."""
     profile = Profile(profile_name, profile_dir)
     if not profile.exists():
         error(f"profile {profile.name} not found at {profile.root}")
         sys.exit(1)
-    profile.root / "config" / "config.py"
+    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)
 
 
 def then_launch(

qbpm/operations.py

@@ -1,12 +1,14 @@
 import shutil
 import subprocess
+from dataclasses import dataclass
 from pathlib import Path
 from sys import platform
 from typing import Optional
 
 from xdg import BaseDirectory
 
 from . import Profile, profiles
+from .icons import find_icon_file
 from .utils import AUTO_MENUS, error, installed_menus
 
 
@@ -64,7 +66,11 @@ def desktop(profile: Profile) -> bool:
 
 
 def choose(
-    profile_dir: Path, menu: str, foreground: bool, qb_args: tuple[str, ...]
+    profile_dir: Path,
+    menu: str,
+    foreground: bool,
+    qb_args: tuple[str, ...],
+    force_icon: bool,
 ) -> bool:
     menu = menu or next(installed_menus())
     if not menu:
@@ -73,19 +79,24 @@ def choose(
     if menu == "applescript" and platform != "darwin":
         error(f"Menu applescript cannot be used on a {platform} host")
         return False
-    profiles = [profile.name for profile in sorted(profile_dir.iterdir())]
+    profiles = [
+        Profile(profile.name, profile_dir) for profile in sorted(profile_dir.iterdir())
+    ]
     if len(profiles) == 0:
         error("No profiles")
         return False
 
-    command = menu_command(menu, profiles, qb_args)
-    if not command:
+    menu_cmd = menu_command(menu, profiles, qb_args)
+    if not menu_cmd:
         return False
 
+    menu_items = build_menu_items(profiles, menu_cmd.icon_support or force_icon)
+
     selection_cmd = subprocess.run(
-        command[1],
-        input=command[0],
+        menu_cmd.command,
+        input=menu_items,
         text=True,
+        # TODO remove shell dependency
         shell=True,
         stdout=subprocess.PIPE,
         stderr=None,
@@ -102,24 +113,33 @@ def choose(
         return False
 
 
+@dataclass
+class Menu:
+    command: str
+    icon_support: bool
+
+
 def menu_command(
-    menu: str, profiles: list[str], qb_args: tuple[str, ...]
-) -> Optional[tuple[str | None, str]]:
+    menu: str, profiles: list[Profile], qb_args: tuple[str, ...]
+) -> Optional[Menu]:
+    icon = False
     arg_string = " ".join(qb_args)
     if menu == "applescript":
-        profile_list = '", "'.join(profiles)
-        return (
-            None,
+        profile_list = '", "'.join([p.name for p in profiles])
+        return Menu(
             f"""osascript -e \'set profiles to {{"{profile_list}"}}
 set profile to choose from list profiles with prompt "qutebrowser: {arg_string}" default items {{item 1 of profiles}}
 item 1 of profile\'""",
+            icon,
         )
 
     prompt = "-p qutebrowser"
     command = menu
+    # TODO arg
     if len(menu.split(" ")) == 1:
         program = Path(menu).name
         if program == "rofi":
+            icon = True
             command = f"{menu} -dmenu -no-custom {prompt} -mesg '{arg_string}'"
         elif program == "wofi":
             command = f"{menu} --dmenu {prompt}"
@@ -128,12 +148,22 @@ def menu_command(
         elif program.startswith("fzf"):
             command = f"{menu} --prompt 'qutebrowser '"
         elif program == "fuzzel":
+            icon = True
             command = f"{menu} -d"
-            profiles = [f"{p}\0icon\x1fqbpm-{p}\n" for p in profiles]
-            print(f"{profiles=}")
     exe = command.split(" ")[0]
     if not shutil.which(exe):
         error(f"command '{exe}' not found")
         return None
-    profile_list = "".join(profiles)
-    return (profile_list, command)
+    return Menu(command, icon)
+
+
+def build_menu_items(profiles: list[Profile], icon: bool) -> str:
+    if icon and any(icons := [find_icon_file(p) for p in profiles]):
+        menu_items = [
+            f"{p.name}\0icon\x1f{icon or 'qutebrowser'}"
+            for (p, icon) in zip(profiles, icons)
+        ]
+    else:
+        menu_items = [p.name for p in profiles]
+
+    return "\n".join(menu_items)

qbpm/profiles.py

@@ -1,16 +1,12 @@
-import subprocess
 from functools import partial
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Optional
 
-import favicon
-import requests
-from PIL import Image
 from xdg import BaseDirectory
 from xdg.DesktopEntry import DesktopEntry
 
 from . import Profile
+from .icons import download_icon
 from .utils import error, user_config_dir
 
 
@@ -48,7 +44,7 @@ def create_config(
 application_dir = Path(BaseDirectory.xdg_data_home) / "applications" / "qbpm"
 
 
-def create_desktop_file(profile: Profile, icon: str | None) -> None:
+def create_desktop_file(profile: Profile, icon: Path | str | None = 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
@@ -60,36 +56,6 @@ def create_desktop_file(profile: Profile, icon: str | None) -> None:
     desktop.write()
 
 
-def download_icon(profile: Profile, home_page: str) -> str | None:
-    icons = favicon.get(home_page)
-    wanted = None
-    print(icons)
-    for icon in icons:
-        if ".ico" in icon.url:
-            wanted = icon
-            print(f"chose {icon}")
-
-    tmp_dir = TemporaryDirectory()
-    work_dir = Path(tmp_dir.name)
-    fav = work_dir / "favicon.ico"
-    if wanted:
-        resp = requests.get(icon.url, timeout=10)
-        with open(fav, "wb") as im:
-            for chunk in resp.iter_content(1024):
-                im.write(chunk)
-
-    name = f"qbpm-{profile.name}"
-    image = Image.open(fav)
-    png = work_dir / "icon.png"
-    image.save(png)
-    subprocess.run(
-        f"xdg-icon-resource install --context apps {png} --size {image.size[0]} {name}",
-        shell=True,
-    )
-    # image.save(Path(BaseDirectory.save_data_path("icons")) / "hicolor" / f"{image.size[0]}x{image.size[0]}" / "apps" / f"{name}.png")
-    return name
-
-
 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")
@@ -111,6 +77,7 @@ def new_profile(
     if create_profile(profile, overwrite):
         create_config(profile, home_page, overwrite)
         if home_page:
+            # TODO catch errors?
             icon = download_icon(profile, home_page)
         if desktop_file:
             create_desktop_file(profile, icon)