From ca4bb4a271f8c03105c21083755a3b870fb46a4d Mon Sep 17 00:00:00 2001
From: Alan Hamlett <alan.hamlett@gmail.com>
Date: Sun, 11 Mar 2018 15:35:35 -0700
Subject: [PATCH] Support for Swift and Objective-C dependencies

---
 tests/samples/codefiles/objective-c.m | 18 ++++++
 tests/samples/codefiles/swift.swift   | 16 +++++
 tests/test_dependencies.py            | 23 ++++++++
 wakatime/dependencies/objective.py    | 84 +++++++++++++++++++++++++++
 4 files changed, 141 insertions(+)
 create mode 100644 tests/samples/codefiles/objective-c.m
 create mode 100644 tests/samples/codefiles/swift.swift
 create mode 100644 wakatime/dependencies/objective.py

diff --git a/tests/samples/codefiles/objective-c.m b/tests/samples/codefiles/objective-c.m
new file mode 100644
index 0000000..f7ebdcd
--- /dev/null
+++ b/tests/samples/codefiles/objective-c.m
@@ -0,0 +1,18 @@
+//
+//  Objective-C.m
+//
+
+#import "SomeViewController.h"
+#import <UIKit.h>
+#import <PromiseKit/Something.h>
+
+@interface ViewController : UIViewController
+  @property (nonnull, strong) URL *url;
+@end
+
+@implementation ViewController {
+- (void)viewDidLoad {
+  [super viewDidLoad];
+  // noop
+}
+@end
diff --git a/tests/samples/codefiles/swift.swift b/tests/samples/codefiles/swift.swift
new file mode 100644
index 0000000..9819118
--- /dev/null
+++ b/tests/samples/codefiles/swift.swift
@@ -0,0 +1,16 @@
+//
+//  Swift.swift
+//
+
+import Foundation
+import UIKit
+import PromiseKit
+
+class ViewController: UIViewController {
+
+  private var attribute: URL?
+
+  override func viewDidLoad() {
+    // noop
+  }
+}
diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py
index 2ad2ef5..88772ae 100644
--- a/tests/test_dependencies.py
+++ b/tests/test_dependencies.py
@@ -397,3 +397,26 @@ class DependenciesTestCase(TestCase):
             expected_lines=37,
             entity='typescript.ts',
         )
+
+    def test_swift_dependencies_detected(self):
+        self.shared(
+            expected_dependencies=[
+                'UIKit',
+                'PromiseKit',
+            ],
+            expected_language='Swift',
+            expected_lines=16,
+            entity='swift.swift',
+        )
+
+    def test_objective_c_dependencies_detected(self):
+        self.shared(
+            expected_dependencies=[
+                'SomeViewController',
+                'UIKit',
+                'PromiseKit',
+            ],
+            expected_language='Objective-C',
+            expected_lines=18,
+            entity='objective-c.m',
+        )
diff --git a/wakatime/dependencies/objective.py b/wakatime/dependencies/objective.py
new file mode 100644
index 0000000..bc06bdf
--- /dev/null
+++ b/wakatime/dependencies/objective.py
@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+"""
+    wakatime.dependencies.objective
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    Parse dependencies from Objective-C and Swift code.
+
+    :copyright: (c) 2018 Alan Hamlett.
+    :license: BSD, see LICENSE for more details.
+"""
+
+import re
+
+from . import TokenParser
+
+
+class SwiftParser(TokenParser):
+    state = None
+    exclude = [
+        r'^foundation$',
+    ]
+
+    def parse(self):
+        for index, token, content in self.tokens:
+            self._process_token(token, content)
+        return self.dependencies
+
+    def _process_token(self, token, content):
+        if self.partial(token) == 'Declaration':
+            self._process_declaration(token, content)
+        elif self.partial(token) == 'Class':
+            self._process_class(token, content)
+        else:
+            self._process_other(token, content)
+
+    def _process_declaration(self, token, content):
+        if self.state is None:
+            self.state = content
+
+    def _process_class(self, token, content):
+        if self.state == 'import':
+            self.append(content)
+        self.state = None
+
+    def _process_other(self, token, content):
+        pass
+
+
+class ObjectiveCParser(TokenParser):
+    state = None
+    extension = re.compile(r'\.[mh]$')
+
+    def parse(self):
+        for index, token, content in self.tokens:
+            self._process_token(token, content)
+        return self.dependencies
+
+    def _process_token(self, token, content):
+        if self.partial(token) == 'Preproc':
+            self._process_preproc(token, content)
+        else:
+            self._process_other(token, content)
+
+    def _process_preproc(self, token, content):
+        if self.state:
+            self._process_import(token, content)
+
+        self.state = content
+
+    def _process_import(self, token, content):
+        if self.state == '#' and content.startswith('import '):
+            self.append(self._format(content))
+        self.state = None
+
+    def _process_other(self, token, content):
+        pass
+
+    def _format(self, content):
+        content = content.strip().lstrip('import ').strip()
+        content = content.strip('"').strip("'").strip()
+        content = content.strip('<').strip('>').strip()
+        content = content.split('/')[0]
+        content = self.extension.sub('', content, count=1)
+        return content