diff --git a/bumblebee_status/modules/contrib/wlrotation.py b/bumblebee_status/modules/contrib/wlrotation.py
index 54b43a9..d94c280 100644
--- a/bumblebee_status/modules/contrib/wlrotation.py
+++ b/bumblebee_status/modules/contrib/wlrotation.py
@@ -7,38 +7,33 @@ Requires the following executable:
     * swaymsg
 """
 
-import bumblebee.util
-import bumblebee.input
-import bumblebee.output
-import bumblebee.engine
+import core.module
+import core.input
+import util.cli
+
 import json
 from os import environ, path
 
 possible_orientations = ["normal", "90", "180", "270"]
 
-class Module(bumblebee.engine.Module):
-    def __init__(self, engine, config):
-        self.service = "sway-auto-rotate@%s" % path.basename(environ['SWAYSOCK']).replace('-', '\\\\x2d')
-        widgets = []
-        self._engine = engine
-        super(Module, self).__init__(engine, config, widgets)
-        self.update_widgets(widgets)
+class Module(core.module.Module):
+    def __init__(self, config, theme):
+        super().__init__(config, theme, [])
 
-    def update_widgets(self, widgets):
-        for display in json.loads(bumblebee.util.execute("/usr/bin/swaymsg -t get_outputs -r")):
+        self.service = "sway-auto-rotate@%s" % path.basename(environ['SWAYSOCK']).replace('-', '\\\\x2d')
+#        self._widgets = []
+        self.update_widgets()
+
+    def update_widgets(self):
+        for display in json.loads(util.cli.execute("/usr/bin/swaymsg -t get_outputs -r")):
             name = display['name']
             widget = self.widget(name)
             if not widget:
-                widget = bumblebee.output.Widget(name=name, full_text="若 "+name)
-                widget.set("auto", bumblebee.util.execute("systemctl --user show %s -p ActiveState" % self.service).split("\n")[0].split("=")[1])
-                self._engine.input.register_callback(widget, button=bumblebee.input.LEFT_MOUSE, cmd=self._toggle)
-                self._engine.input.register_callback(widget, button=bumblebee.input.RIGHT_MOUSE, cmd=self._toggle_auto)
+                widget = self.add_widget(name=name, full_text="若 "+name)
+                widget.set("auto", util.cli.execute("systemctl --user show %s -p ActiveState" % self.service).split("\n")[0].split("=")[1])
+                core.input.register(widget, button=core.input.LEFT_MOUSE, cmd=self._toggle)
+                core.input.register(widget, button=core.input.RIGHT_MOUSE, cmd=self._toggle_auto)
             widget.set("orientation", display['transform'])
-            widgets.append(widget)
-
-    def update(self, widgets):
-        if len(widgets) <= 0:
-            self.update_widgets(widgets)
 
     def state(self, widget):
         curr = widget.get("auto", "inactive")
@@ -47,23 +42,22 @@ class Module(bumblebee.engine.Module):
         return [curr]
 
     def _toggle_auto(self, event):
-        widget = self.widget_by_id(event["instance"])
+        widget = self.widget(widget_id=event["instance"])
         if widget.get("auto") == "inactive":
-            bumblebee.util.execute("systemctl --user start %s" % self.service)
+            util.cli.execute("systemctl --user start %s" % self.service)
         else:
-            bumblebee.util.execute("systemctl --user stop %s" % self.service)
-        widget.set("auto", bumblebee.util.execute("systemctl --user show %s -p ActiveState" % self.service).split("\n")[0].split("=")[1])
+            util.cli.execute("systemctl --user stop %s" % self.service)
+        widget.set("auto", util.cli.execute("systemctl --user show %s -p ActiveState" % self.service).split("\n")[0].split("=")[1])
 
     def _toggle(self, event):
-        widget = self.widget_by_id(event["instance"])
+        widget = self.widget(widget_id=event["instance"])
 
         # compute new orientation based on current orientation
-        idx = possible_orientations.index(widget.get("orientation"))
-        idx = (idx + 1) % len(possible_orientations)
+        idx = (possible_orientations.index(widget.get("orientation"))+ 1) % len(possible_orientations)
         new_orientation = possible_orientations[idx]
 
         widget.set("orientation", new_orientation)
 
-        bumblebee.util.execute("swaymsg 'output {} transform {}'".format(widget.name, new_orientation))
+        util.cli.execute("swaymsg 'output {} transform {}'".format(widget.name, new_orientation))
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4