wallpaper: refactor a bit wallpaper module

bin/wallpaper

@@ -1,77 +1,72 @@
 #!/usr/bin/env python3
 
-# Build a multi screen wallpaper
+"""Build a multi screen wallpaper."""
-
-# First argument is the directory where the wallpapers can be
-# found. We use xinerama to know the dimension of each screen.
 
 import os
+import sys
 import random
 import argparse
 import tempfile
 import collections
 import itertools
+import logging
+import logging.handlers
 
 from Xlib import display
 from Xlib.ext import randr
-
 from PIL import Image
+from systemd import journal
 
-parser = argparse.ArgumentParser()
-parser.add_argument(
-    "-d",
-    "--directory",
-    default=".",
-    help="search for images in DIRECTORY",
-)
-parser.add_argument(
-    "-t",
-    "--target",
-    default="background.png",
-    help="write background to FILE",
-)
-parser.add_argument(
-    "--extra-images",
-    default=3,
-    metavar="N",
-    help="consider N additional images to choose the best combination",
-)
-parser.add_argument(
-    "--compression", default=0, type=int, help="compression level when saving"
-)
-options = parser.parse_args()
-background = None
 
-# Get display size
+logger = logging.getLogger("wallpaper")
-d = display.Display()
+Rectangle = collections.namedtuple("Rectangle", ["x", "y", "width", "height"])
-screen = d.screen()
+WallpaperPart = collections.namedtuple("WallpaperPart", ["rectangle", "image"])
-window = screen.root.create_window(0, 0, 1, 1, 1, screen.root_depth)
-background = Image.new("RGB", (screen.width_in_pixels, screen.height_in_pixels))
 
-# Query randr extension
+def get_outputs():
-Output = collections.namedtuple("Output", ["x", "y", "width", "height"])
+    """Get physical outputs."""
-outputs = []
+    # Get display size
-screen_resources = randr.get_screen_resources_current(window)
+    d = display.Display()
-for output in screen_resources.outputs:
+    screen = d.screen()
+    window = screen.root.create_window(0, 0, 1, 1, 1, screen.root_depth)
+    background = Image.new("RGB", (screen.width_in_pixels, screen.height_in_pixels))
+
+    # Query randr extension
+    outputs = []
+    screen_resources = randr.get_screen_resources_current(window)
+    for output in screen_resources.outputs:
         output_info = randr.get_output_info(window, output, screen_resources.timestamp)
         if output_info.crtc == 0:
             continue
         crtc_info = randr.get_crtc_info(window, output_info.crtc, output_info.timestamp)
-    outputs.append(Output(crtc_info.x, crtc_info.y, crtc_info.width, crtc_info.height))
+        outputs.append(Rectangle(crtc_info.x, crtc_info.y, crtc_info.width, crtc_info.height))
-if not outputs:
-    outputs.append(Output(0, 0, (background.width, background.height)))
 
-# Get all possible rectangles from outputs
+    for o in outputs:
-candidates = []
+        logger.debug("output: %s", o)
-for output in outputs:
+    return outputs, background
-    candidates.append(output)
+
+
+def get_covering_rectangles(outputs):
+    """Compute all possible groups of covering boxes for the provided
+    outputs. For each group, an output is included in exactly one box.
+
+    >>> get_covering_rectangles([Rectangle(0, 0, 100, 100)])
+    {(Rectangle(x=0, y=0, width=100, height=100),)}
+    >>> get_covering_rectangles([Rectangle(0, 0, 100, 100), Rectangle(100, 0, 100, 100)])
+    {(Rectangle(x=0, y=0, width=100, height=100), Rectangle(x=100, y=0, width=100, height=100)), (Rectangle(x=0, y=0, width=200, height=100),)}
+    >>> get_covering_rectangles([Rectangle(0, 0, 100, 100), Rectangle(100, 0, 100, 100), Rectangle(0, 100, 100, 100)])
+    {(Rectangle(x=100, y=0, width=100, height=100), Rectangle(x=0, y=100, width=100, height=100), Rectangle(x=0, y=0, width=100, height=100)), (Rectangle(x=100, y=0, width=100, height=100), Rectangle(x=0, y=0, width=100, height=200)), (Rectangle(x=0, y=0, width=200, height=100), Rectangle(x=0, y=100, width=100, height=100))}
+    """
+
+    candidates = set()
+    for output in outputs:
+        candidates.add(output)
         for ooutput in outputs:
             if ooutput == output:
                 continue
             if output.x > ooutput.x or output.y > ooutput.y:
                 continue
-        candidates.append(
+            candidates.add(
-            Output(
+                Rectangle(
                     output.x,
                     output.y,
                     ooutput.x - output.x + ooutput.width,
@@ -79,9 +74,9 @@ for output in outputs:
                 )
             )
 
-# Get all rectangle combinations to cover outputs without overlapping
+    # Get all rectangle combinations to cover outputs without overlapping
-groups = []
+    groups = set()
-for r in range(len(candidates)):
+    for r in range(len(candidates)):
         for candidate in itertools.combinations(candidates, r + 1):
             for output in outputs:
                 nb = 0
@@ -91,79 +86,145 @@ for r in range(len(candidates)):
                 if nb != 1:  # output not contained in a single rectangle
                     break
             else:
-            groups.append(candidate)
+                groups.add(candidate)
 
-# Get random images
+    for g in groups:
-images = []
+        logger.debug("group: %s", g)
-for base, _, files in os.walk(os.path.join(options.directory)):
+    return groups
+
+def get_random_images(directory, number):
+    """Get random images from a directory."""
+    images = []
+    for base, _, files in os.walk(os.path.join(directory)):
         for i in files:
             if os.path.splitext(i)[1].lower() in (".jpg", ".jpeg", ".png", ".webp"):
                 images.append(os.path.join(base, i))
-images = random.sample(images, len(outputs) * options.extra_images)
+    images = random.sample(images, number)
-images = [Image.open(image) for image in images]
+    images = [Image.open(image) for image in images]
 
-# Find optimal combination for images
+    for image in images:
-best_association = None
+        logger.debug("image: %s %s×%s", image.filename[(len(directory) + 1) :], *image.size)
-best_score = 0
+    return images
-for group in groups:
+
+def get_best_parts(groups, images):
+    """Find optimal association for images for the groups of covering rectangles."""
+    best_association = None
+    best_score = 0
+    for group in groups:
         for association in (zip(group, p) for p in itertools.permutations(images)):
             score = 0
-        association = list(association)
+            association = [WallpaperPart(rectangle=assoc[0], image=assoc[1]) for assoc in association]
-        for output, image in association:
+            for assoc in association:
                 # Similar ratio
-            oratio = output.width * 100 // output.height
+                oratio = assoc.rectangle.width * 100 // assoc.rectangle.height
-            iratio = image.width * 100 // image.height
+                iratio = assoc.image.width * 100 // assoc.image.height
                 r = iratio / oratio
                 if r > 1:
                     r = 1 / r
                 score += r * 100
                 # Similar scale (when cropped)
-            opixels = output.width * output.height
+                opixels = assoc.rectangle.width * assoc.rectangle.height
-            ipixels = image.width * image.height * r
+                ipixels = assoc.image.width * assoc.image.height * r
                 r = ipixels / opixels
                 if r >= 1:
                     r = 1
                 score += r * 50
             score /= len(group)
+            logger.debug("association: %s, score %.2f", association, score)
             if score > best_score or best_association is None:
                 best_association = association
                 best_score = score
 
-print(
+    return best_association
-    "wallpaper: {}".format(
-        " + ".join(
-            "`{}` ({}×{})".format(
-                image.filename[(len(options.directory) + 1) :], *image.size
-            )
-            for image in [couple[1] for couple in best_association]
-        )
-    )
-)
 
-for couple in best_association:
-    output = couple[0]
-    image = couple[1]
 
-    # Find the right size for the screen
+def build(background, wallpaper_parts):
-    imx, imy = output.width, image.height * output.width // image.width
+    """Stitch wallpaper into provided background."""
-    if imy < output.height:
+    for part in wallpaper_parts:
-        imx, imy = image.width * output.height // image.height, output.height
+        rectangle = part.rectangle
+        image = part.image
+
+        imx, imy = rectangle.width, image.height * rectangle.width // image.width
+        if imy < rectangle.height:
+            imx, imy = image.width * rectangle.height // image.height, rectangle.height
         if image.size != (imx, imy):
             image = image.resize((imx, imy), Image.LANCZOS)
         image = image.crop(
             (
-            (imx - output.width) / 2,
+                (imx - rectangle.width) / 2,
-            (imy - output.height) / 2,
+                (imy - rectangle.height) / 2,
-            imx - (imx - output.width) / 2,
+                imx - (imx - rectangle.width) / 2,
-            imy - (imy - output.height) / 2,
+                imy - (imy - rectangle.height) / 2,
             )
         )
-    background.paste(image, (output.x, output.y))
+        background.paste(image, (rectangle.x, rectangle.y))
 
-# Save
+
-assert background, "Don't know the size of the display area"
+def save(wallpaper, target, compression):
-with tempfile.NamedTemporaryFile(
+    """Save wallpaper to target."""
-    delete=False, dir=os.path.dirname(os.path.realpath(options.target))
+    with tempfile.NamedTemporaryFile(
-) as tmp:
+        delete=False, dir=os.path.dirname(os.path.realpath(target))
-    background.save(tmp, "png", compress_level=options.compression)
+    ) as tmp:
-    os.rename(tmp.name, options.target)
+        background.save(tmp, "png", compress_level=compression)
+        os.rename(tmp.name, target)
+
+if __name__ == "__main__":
+    # Parse
+    description = sys.modules[__name__].__doc__
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        default=False,
+        help="enable debugging",
+    )
+    parser.add_argument(
+        "-d",
+        "--directory",
+        default=".",
+        help="search for images in DIRECTORY",
+    )
+    parser.add_argument(
+        "-t",
+        "--target",
+        default="background.png",
+        help="write background to FILE",
+    )
+    parser.add_argument(
+        "--extra-images",
+        default=3,
+        metavar="N",
+        help="consider N additional images to choose the best combination",
+    )
+    parser.add_argument(
+        "--compression", default=0, type=int, help="compression level when saving"
+    )
+    options = parser.parse_args()
+
+    # Logging
+    root = logging.getLogger("")
+    root.setLevel(logging.WARNING)
+    logger.setLevel(options.debug and logging.DEBUG or logging.INFO)
+    if sys.stderr.isatty():
+        ch = logging.StreamHandler()
+        ch.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
+        root.addHandler(ch)
+    else:
+        root.addHandler(journal.JournalHandler(SYSLOG_IDENTIFIER=logger.name))
+
+    try:
+        outputs, background = get_outputs()
+        candidates = get_covering_rectangles(outputs)
+        images = get_random_images(options.directory, len(outputs) + options.extra_images)
+        wallpaper_parts = get_best_parts(candidates, images)
+        for part in wallpaper_parts:
+            logger.info("wallpaper: {} ({}×{})".format(
+                part.image.filename[(len(options.directory) + 1) :],
+                *part.image.size
+            ))
+        build(background, wallpaper_parts)
+        save(background, options.target, options.compression)
+    except Exception as e:
+        logger.exception("%s", e)
+        sys.exit(1)
+    sys.exit(0)