diff --git a/modules/theme/config.nix b/modules/theme/config.nix
index 01be14c..f6544b3 100644
--- a/modules/theme/config.nix
+++ b/modules/theme/config.nix
@@ -1,9 +1,5 @@
-{
-  config,
-  lib,
-  ...
-}: let
-  inherit (lib) mkDefault;
+{lib, ...}: let
+  inherit (lib.modules) mkDefault;
 in {
   config = {
     vim.theme = {
diff --git a/modules/theme/default.nix b/modules/theme/default.nix
index 02f4f37..3b37928 100644
--- a/modules/theme/default.nix
+++ b/modules/theme/default.nix
@@ -1,4 +1,4 @@
-{...}: {
+{
   imports = [
     ./theme.nix
     ./config.nix
diff --git a/modules/theme/theme.nix b/modules/theme/theme.nix
index 094a56c..2ae3e45 100644
--- a/modules/theme/theme.nix
+++ b/modules/theme/theme.nix
@@ -3,44 +3,50 @@
   lib,
   ...
 }: let
-  inherit (lib) mkOption types attrNames mkIf nvim;
+  inherit (lib.options) mkOption;
+  inherit (lib.attrsets) attrNames;
+  inherit (lib.types) bool lines enum;
+  inherit (lib.modules) mkIf;
+  inherit (lib.nvim.dag) entryBefore;
 
   cfg = config.vim.theme;
   supported_themes = import ./supported_themes.nix {inherit lib;};
 in {
   options.vim.theme = {
     enable = mkOption {
-      type = types.bool;
+      type = bool;
       description = "Enable theming";
     };
 
     name = mkOption {
-      type = types.enum (attrNames supported_themes);
+      type = enum (attrNames supported_themes);
       description = "Supported themes can be found in `supported_themes.nix`";
     };
 
     style = mkOption {
-      type = with types; enum supported_themes.${cfg.name}.styles;
+      type = enum supported_themes.${cfg.name}.styles;
       description = "Specific style for theme if it supports it";
     };
 
     transparent = mkOption {
-      type = with types; bool;
+      type = bool;
       default = false;
       description = "Whether or not transparency should be enabled. Has no effect for themes that do not support transparency";
     };
 
     extraConfig = mkOption {
-      type = with types; lines;
+      type = lines;
       description = "Additional lua configuration to add before setup";
     };
   };
 
   config = mkIf cfg.enable {
-    vim.startPlugins = [cfg.name];
-    vim.luaConfigRC.themeSetup = nvim.dag.entryBefore ["theme"] cfg.extraConfig;
-    vim.luaConfigRC.theme = supported_themes.${cfg.name}.setup (with cfg; {
-      inherit style transparent;
-    });
+    vim = {
+      startPlugins = [cfg.name];
+      luaConfigRC = {
+        themeSetup = entryBefore ["theme"] cfg.extraConfig;
+        theme = supported_themes.${cfg.name}.setup (with cfg; {inherit style transparent;});
+      };
+    };
   };
 }