Check element name when generating storyboard get viewController func…

…tions Fixes #11
mac-cain13 · Feb 1, 2015 · b8bce81 · b8bce81
1 parent 67b1681
commit b8bce81
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/R.swift/types.swift b/R.swift/types.swift
@@ -18,6 +18,7 @@ struct Type: Printable {
   static let _UINib = Type(className: "UINib")
   static let _UIImage = Type(className: "UIImage")
   static let _UIStoryboard = Type(className: "UIStoryboard")
+  static let _UIViewController = Type(className: "UIViewController")
 
   let moduleName: String?
   let className: String
@@ -218,19 +219,21 @@ class StoryboardParserDelegate: NSObject, NSXMLParserDelegate {
       }
 
     default:
-      if let viewController = viewControllerFromAttributes(attributeDict) {
+      if let viewController = viewControllerFromAttributes(attributeDict, elementName: elementName) {
         viewControllers.append(viewController)
       }
     }
   }
 
-  func viewControllerFromAttributes(attributeDict: [NSObject : AnyObject]) -> Storyboard.ViewController? {
+  func viewControllerFromAttributes(attributeDict: [NSObject : AnyObject], elementName: String) -> Storyboard.ViewController? {
     if attributeDict["sceneMemberID"] as? String == "viewController" {
       if let storyboardIdentifier = attributeDict["storyboardIdentifier"] as? String {
         let customModule = attributeDict["customModule"] as? String
-        let customClass = attributeDict["customClass"] as? String ?? "UIViewController"
+        let customClass = attributeDict["customClass"] as? String
+        let customType = customClass.map { Type(moduleName: customModule, className: $0, optional: false) }
 
-        return Storyboard.ViewController(storyboardIdentifier: storyboardIdentifier, type: Type(moduleName: customModule, className: customClass, optional: false))
+        let type = customType ?? ElementNameToTypeMapping[elementName] ?? Type._UIViewController
+        return Storyboard.ViewController(storyboardIdentifier: storyboardIdentifier, type: type)
       }
     }
 

diff --git a/R.swift/values.swift b/R.swift/values.swift
@@ -42,3 +42,15 @@ let Ordinals = [
   (number: 19, word: "nineteenth"),
   (number: 20, word: "twentieth"),
 ]
+
+let ElementNameToTypeMapping = [
+  "viewController": Type._UIViewController,
+  "glkViewController": Type(className: "GLKViewController"),
+  "tabBarController": Type(className: "UITabBarController"),
+  "pageViewController": Type(className: "UIPageViewController"),
+  "tableViewController": Type(className: "UITableViewController"),
+  "splitViewController": Type(className: "UISplitViewController"),
+  "navigationController": Type(className: "UINavigationController"),
+  "avPlayerViewController": Type(className: "AVPlayerViewController"),
+  "collectionViewController": Type(className: "UICollectionViewController"),
+]