HTMLReader/Code/HTMLDocument.m at master · yanflex/HTMLReader · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
//  HTMLDocument.m
//
//  Public domain. https://github.com/nolanw/HTMLReader

#import "HTMLDocument.h"
#import "HTMLParser.h"

@implementation HTMLDocument

+ (instancetype)documentWithString:(NSString *)string
{
    HTMLParser *parser = [[HTMLParser alloc] initWithString:string context:nil];
    return parser.document;
}

- (instancetype)initWithString:(NSString *)string
{
    return [self.class documentWithString:string];
}

- (HTMLDocumentType *)documentType
{
    return FirstNodeOfType(self.children, [HTMLDocumentType class]);
}

- (void)setDocumentType:(HTMLDocumentType *)documentType
{
    HTMLDocumentType *oldDocumentType = self.documentType;
    NSMutableOrderedSet *children = [self mutableChildren];
    if (oldDocumentType && documentType) {
        NSUInteger i = [children indexOfObject:oldDocumentType];
        [children replaceObjectAtIndex:i withObject:documentType];
    } else if (documentType) {
        HTMLElement *rootElement = self.rootElement;
        if (rootElement) {
            [children insertObject:documentType atIndex:[children indexOfObject:rootElement]];
        } else {
            [children addObject:documentType];
        }
    } else if (oldDocumentType) {
        [children removeObject:oldDocumentType];
    }
}

- (HTMLElement *)rootElement
{
    return FirstNodeOfType(self.children, [HTMLElement class]);
}

- (void)setRootElement:(HTMLElement *)rootElement
{
    HTMLElement *oldRootElement = self.rootElement;
    NSMutableOrderedSet *children = [self mutableChildren];
    if (oldRootElement && rootElement) {
        [children replaceObjectAtIndex:[children indexOfObject:oldRootElement] withObject:rootElement];
    } else if (rootElement) {
        [children addObject:rootElement];
    } else if (oldRootElement) {
        [children removeObject:oldRootElement];
    }
}

static id FirstNodeOfType(id <NSFastEnumeration> collection, Class type)
{
    for (id node in collection) {
        if ([node isKindOfClass:type]) {
            return node;
        }
    }
    return nil;
}

@end