add color in table and token parameter

1 files changed, 7 insertions, 3 deletions

diff --git a/pgcli/pgstyle.py b/pgcli/pgstyle.py
index 2ff92ed1..0a6ac393 100644
--- a/pgcli/pgstyle.py
+++ b/pgcli/pgstyle.py
@@ -10,8 +10,12 @@ def style_factory(name, cli_style):
     except ClassNotFound:
         style = pygments.styles.get_style_by_name('native')
 
-    custom_styles = dict([(string_to_tokentype(x), y)
-                            for x, y in cli_style.items()])
-
+    custom_styles = {}
+    for token in cli_style:
+        try:
+            custom_styles[string_to_tokentype(
+                token)] = style.styles[string_to_tokentype(cli_style[token])]
+        except AttributeError as err:
+            custom_styles[string_to_tokentype(token)] = cli_style[token]
     return PygmentsStyle.from_defaults(style_dict=custom_styles,
                                        pygments_style_cls=style)