diff --git a/src/main.cpp b/src/main.cpp
index c2a3bc2..3602805 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -53,27 +53,32 @@ int main()
     while (std::getline(std::cin, line))
     {
         const auto line_nc{remove_colors(line)};
-        // Only look at lines starting with + or -
+        // Only look at long enough lines starting with + or -
         if (line_nc.length() >= 2 && (line_nc[0] == '+' || line_nc[0] == '-'))
         {
             // Don't color +++ and --- lines
             if (line_nc.length() < 3
                 || (line_nc[2] != '+' && line_nc[2] != '-'))
             {
-                if (line_nc[1] == '+')
+                const auto color{[&line_nc]() -> std::string_view
+                                 {
+                                     if (line_nc[1] == '+')
+                                     {
+                                         return green;
+                                     }
+                                     if (line_nc[1] == '-')
+                                     {
+                                         return red;
+                                     }
+                                     return {};
+                                 }()};
+                if (!color.empty())
                 {
                     const auto pos{get_start_pos(line)};
-                    std::cout << line.substr(0, pos) << green
+                    std::cout << line.substr(0, pos) << color
                               << line.substr(pos) << reset << '\n';
                     continue;
                 }
-                if (line_nc[1] == '-')
-                {
-                    const auto pos{get_start_pos(line)};
-                    std::cout << line.substr(0, pos) << red << line.substr(pos)
-                              << reset << '\n';
-                    continue;
-                }
             }
         }
         std::cout << line << '\n';