Merge pull request #1 from frostming/fix/command

feat: use subclass instead of monkeypatch

hy_tg.py

@@ -1,58 +1,54 @@
 import argparse
-import random
-import os
 import gc
+import io
+import random
 import subprocess
-from telebot import TeleBot  # type: ignore
-from telebot.types import Message  # type: ignore
 
-from prettymapp.geo import get_aoi
-from prettymapp.osm import get_osm_geometries
-from prettymapp.plotting import Plot
-from prettymapp.settings import STYLES
-
-from matplotlib import figure
-
-from PIL import Image
 import numpy as np
 import PIL
-import io
+from matplotlib import figure
+from PIL import Image
+from prettymapp.geo import get_aoi
+from prettymapp.osm import get_osm_geometries
+from prettymapp.plotting import Plot as PrettyPlot
+from prettymapp.settings import STYLES
+from telebot import TeleBot  # type: ignore
+from telebot.types import Message  # type: ignore
 
 PIL.Image.MAX_IMAGE_PIXELS = 933120000
 file_in = "map.jpg"
 file_out = "map_out.jpg"
 
 
-# monkey patch for Plot thanks @higuoxing https://github.com/higuoxing
-def __post_init__(self):
-    (
-        self.xmin,
-        self.ymin,
-        self.xmax,
-        self.ymax,
-    ) = self.aoi_bounds
-    # take from aoi geometry bounds, otherwise probelematic if unequal geometry distribution over plot.
-    self.xmid = (self.xmin + self.xmax) / 2
-    self.ymid = (self.ymin + self.ymax) / 2
-    self.xdif = self.xmax - self.xmin
-    self.ydif = self.ymax - self.ymin
+class Plot(PrettyPlot):
+    # memory leak fix for Plot. thanks @higuoxing https://github.com/higuoxing
+    # refer to: https://www.mail-archive.com/matplotlib-users@lists.sourceforge.net/msg11809.html
+    def __post_init__(self):
+        (
+            self.xmin,
+            self.ymin,
+            self.xmax,
+            self.ymax,
+        ) = self.aoi_bounds
+        # take from aoi geometry bounds, otherwise probelematic if unequal geometry distribution over plot.
+        self.xmid = (self.xmin + self.xmax) / 2
+        self.ymid = (self.ymin + self.ymax) / 2
+        self.xdif = self.xmax - self.xmin
+        self.ydif = self.ymax - self.ymin
 
-    self.bg_buffer_x = (self.bg_buffer / 100) * self.xdif
-    self.bg_buffer_y = (self.bg_buffer / 100) * self.ydif
+        self.bg_buffer_x = (self.bg_buffer / 100) * self.xdif
+        self.bg_buffer_y = (self.bg_buffer / 100) * self.ydif
 
-    # self.fig, self.ax = subplots(
-    #     1, 1, figsize=(12, 12), constrained_layout=True, dpi=1200
-    # )
-    self.fig = figure.Figure(figsize=(12, 12), constrained_layout=True, dpi=1200)
-    self.ax = self.fig.subplots(1, 1)
-    self.ax.set_aspect(1 / np.cos(self.ymid * np.pi / 180))
+        # self.fig, self.ax = subplots(
+        #     1, 1, figsize=(12, 12), constrained_layout=True, dpi=1200
+        # )
+        self.fig = figure.Figure(figsize=(12, 12), constrained_layout=True, dpi=1200)
+        self.ax = self.fig.subplots(1, 1)
+        self.ax.set_aspect(1 / np.cos(self.ymid * np.pi / 180))
 
-    self.ax.axis("off")
-    self.ax.set_xlim(self.xmin - self.bg_buffer_x, self.xmax + self.bg_buffer_x)
-    self.ax.set_ylim(self.ymin - self.bg_buffer_y, self.ymax + self.bg_buffer_y)
-
-
-Plot.__post_init__ = __post_init__
+        self.ax.axis("off")
+        self.ax.set_xlim(self.xmin - self.bg_buffer_x, self.xmax + self.bg_buffer_x)
+        self.ax.set_ylim(self.ymin - self.bg_buffer_y, self.ymax + self.bg_buffer_y)
 
 
 def sizeof_image(image):
@@ -96,14 +92,25 @@ def main():
 
     # Init bot
     bot = TeleBot(options.tg_token)
-    bot_name = bot.get_me().username
-    bot.delete_my_commands(scope=None, language_code=None)
+    bot.set_my_commands(
+        [
+            {
+                "command": "github",
+                "description": "github poster: /github <github_user_name> [<start>-<end>]",
+            },
+            {
+                "command": "map",
+                "description": "pretty map: /map <address>",
+            },
+        ]
+    )
     print("Bot init done.")
 
+    @bot.message_handler(commands=["github"])
     @bot.message_handler(regexp="^github:")
     def github_poster_handler(message: Message):
         reply_message = bot.reply_to(message, "Generating poster please wait:")
-        m = message.text.strip()[7:]
+        m = message.text.strip().split(maxsplit=1)[1].strip()
         message_list = m.split(",")
         name = message_list[0].strip()
         cmd_list = ["github_poster", "github", "--github_user_name", name, "--me", name]
@@ -112,7 +119,7 @@ def main():
             cmd_list.append("--year")
             cmd_list.append(years.strip())
         r = subprocess.check_output(cmd_list).decode("utf-8")
-        if r.find("done") > 0:
+        if "done" in r:
             try:
                 # TODO windows path
                 r = subprocess.check_output(
@@ -127,12 +134,13 @@ def main():
                 bot.reply_to(message, "Something wrong please check")
         bot.delete_message(reply_message.chat.id, reply_message.message_id)
 
+    @bot.message_handler(commands=["map"])
     @bot.message_handler(regexp="^map:")
     def map_handler(message: Message):
         reply_message = bot.reply_to(
             message, "Generating pretty map may take some time please wait:"
         )
-        m = message.text.strip()[4:]
+        m = message.text.strip().split(maxsplit=1)[1].strip()
         location = m.strip()
         styles_list = list(STYLES.keys())
         style = random.choice(styles_list)